"""Native dag-ml run paths for the host backend.
One ``_run_*`` function per supported pipeline shape — single-variant CV+refit, native param-model
generation, repetition (group-aware), rep-fusion reshape, sample-augmentation (global / fold-local),
and the separation / by_source / duplication / stacking branch fan-outs. Each assembles the compat or
canonical DSL, drives dag-ml-cli through the process adapter, and maps the native ``bundle.scores``
into a :class:`~nirs4all.api.result.RunResult`.
"""
from __future__ import annotations
import copy
from collections.abc import Callable
from pathlib import Path
from typing import Any, cast
import numpy as np
from sklearn.base import BaseEstimator, RegressorMixin, TransformerMixin
from nirs4all.api.result import RunResult
from nirs4all.core.metrics import eval_list, get_default_metrics, is_higher_better
from nirs4all.data.ensemble_utils import EnsembleUtils
from nirs4all.data.predictions import Predictions
from nirs4all.pipeline.dagml_bridge import controller_manifests
from .cli_runner import assemble_constrained_cv_refit_dsl, assemble_cv_refit_dsl
from .detect import _generation_kind, _is_augmentation_step, _is_constrained_operator_generator, _is_rep_fusion_step, _is_unconstrained_operator_generator
from .envelope import build_envelope, build_fold_set
from .errors import DagMlUnsupported, _raise_run_failure, _reject_multi_model
from .exclude import FoldLocalExclusion
from .folds import FoldLocalFolds, _build_folds, _build_group_folds, _is_repetition_dataset, _repetition_grain, _split_base_samples, _split_group_grain, _split_pool
from .identity import mint_identity
from .in_process_runner import run_cv_refit_bundle_router as run_cv_refit_bundle
from .methods_lane import HOST_CALLBACK_LANE, lane_record, merge_lane_records, record_execution_lane, run_cv_refit_lane
from .result import _frames_by_variant, _native_variant_config_map, _project_operator_sweep, _scores_to_run_result
from .steps import _apply_model_params, _apply_plain_model_params, _assert_supported_operators, _is_split_step, _model_name, _split_pipeline, _supported_body_steps
[docs]
class DuplicationFusionEstimator(RegressorMixin, DuplicationBranchMergeTransformer):
"""Fit branch-local models on one native training pool and average predictions."""
def __init__(self, branches: list[dict[str, Any]], merge_mode: str = "all") -> None:
super().__init__(branches, merge_mode)
[docs]
def predict(self, X: Any) -> np.ndarray:
x = np.asarray(X)
predictions = [
self._as_2d(branch["model"].predict(self._transform_branch_features(x, branch["transforms"])))
for branch in self._fitted_branches
]
return np.asarray(np.mean(predictions, axis=0))
def _native_param_winner_config_name(
refit_artifacts: list[dict[str, Any]] | None,
variant_config_names: list[str] | None,
variant_model_params: list[dict[str, Any]] | None,
) -> str | None:
"""The WINNING param-sweep variant's legacy ``config_name``, recovered BY CONTENT, else ``None``.
The native run refits the true CV-best variant but emits an opaque variant hash with NO params in the
reports, so the winner's specific ``config_name`` cannot be read off the ScoreSet. The fitted REFIT
estimator IS authoritative, though — it carries the winning param VALUES. We unwrap its bare model
(the last step of the refit sklearn ``Pipeline``, else the estimator itself), then find the variant
whose recorded swept-param values (``variant_model_params[i]``, aligned 1:1 with
``variant_config_names``) all match the winner's — returning ``variant_config_names[i]``. Only the
SWEPT keys are compared (the keys present in the per-variant param dict), so a default param the
estimator also carries does not break the match. Returns ``None`` (→ the positional ``names[0]``
fallback) when there is no refit estimator, no per-variant params, or no unambiguous single match —
never a wrong label.
"""
from sklearn.pipeline import Pipeline
if not refit_artifacts or not variant_config_names or not variant_model_params:
return None
if len(variant_model_params) != len(variant_config_names):
return None
estimator = refit_artifacts[0].get("estimator")
if estimator is None:
return None
model = estimator.steps[-1][1] if isinstance(estimator, Pipeline) and estimator.steps else estimator
if not hasattr(model, "get_params"):
return None
winner_params = model.get_params()
matches = {
config_name
for config_name, variant_params in zip(variant_config_names, variant_model_params, strict=True)
if variant_params and all(key in winner_params and winner_params[key] == value for key, value in variant_params.items())
}
# A single matched config_name keys the winner. Distinct duplicate variants that share the SAME
# swept-param values collapse to ONE config_name (the display hash is the param content), so a grid
# like `scale:[True,False,True]` still resolves unambiguously. Only a genuine multi-NAME match (an
# ill-posed grid where the winner's params match two DIFFERENT names) or no match falls back to the
# positional `names[0]` pairing rather than pick arbitrarily.
return next(iter(matches)) if len(matches) == 1 else None
def _native_param_config_map_from_catalog(
catalog: list[dict[str, Any]],
config_names: list[str] | None,
model_params: list[dict[str, Any]] | None,
) -> dict[str, str]:
"""Match native param variants to legacy config names by their Rust-planned overrides."""
if not config_names or not model_params or len(config_names) != len(model_params):
return {}
matched: dict[str, str] = {}
for variant in catalog:
variant_id = variant.get("variant_id")
choices = variant.get("choices")
if not isinstance(variant_id, str) or not isinstance(choices, dict):
continue
overrides: dict[str, Any] = {}
for choice in choices.values():
if not isinstance(choice, dict):
continue
for override in choice.get("param_overrides", []):
if isinstance(override, dict) and isinstance(override.get("params"), dict):
overrides.update(override["params"])
if not overrides:
continue
names = {
name
for name, params in zip(config_names, model_params, strict=True)
if all(params.get(key) == value for key, value in overrides.items())
}
if len(names) == 1:
matched[variant_id] = names.pop()
return matched
def _run_native_generation(
pipeline: list[Any],
spectro: Any,
dataset_arg: str,
cli: str,
venv_python: str,
run_dir: Path,
metric: str,
task_type: str,
cv_pool: list[int] | None = None,
excluded: set[int] | None = None,
tags_by_sample: dict[int, list[str]] | None = None,
dataset_pickle: str | None = None,
config_name: str = "",
variant_config_names: list[str] | None = None,
variant_model_params: list[dict[str, Any]] | None = None,
random_state: int | None = None,
refit: bool = True,
refit_top_k: int = 1,
) -> RunResult:
"""Run a param-level model sweep as ONE native dag-ml generation + SELECT + refit run.
The model step keeps its generator dict so the bridge lowers it to native ``generators``; we
apply only the plain (non-generator) sibling params to the model, never the sweep. dag-ml
expands the variants, scores EACH by its cross-fold OOF ``metric`` and surfaces every variant's
validation reports (#55), refitting only the winner — ``bundle.scores`` is mapped to a RunResult
as the full PER-VARIANT legacy table (``variant_config_names`` carries the ordered legacy per-variant
config names), so a sweep's num_predictions matches legacy.
``cv_pool`` is the CV sample-int universe (the de-excluded pool in legacy mode, the full train
in opt-in mode); ``excluded`` is marked in the envelope only in opt-in (``keep_in_oof=True``).
"""
steps, splitter = _split_pipeline(pipeline)
if splitter is None:
raise DagMlUnsupported("engine='dag-ml' requires a cross-validator step (e.g. KFold) in the pipeline")
_reject_multi_model(steps)
_assert_supported_operators(steps)
steps = _apply_plain_model_params(steps)
identity = mint_identity(spectro)
pool = list(cv_pool) if cv_pool is not None else spectro.index_column("sample", {"partition": "train"})
folds = _build_folds(splitter, spectro, pool, excluded or set())
envelope = build_envelope(spectro, identity, sample_ints=pool, excluded_sample_ints=excluded or None, tags_by_sample=tags_by_sample, group_by_sample=_split_group_grain(splitter, spectro, pool))
dsl = assemble_cv_refit_dsl(steps, identity, envelope, folds, dsl_id="nirs4all-pipeline", n_splits=len(folds))
def run_callback() -> dict[str, Any]:
import dag_ml
dsl["root_seed"] = random_state if random_state is not None else 0
graph = dag_ml.compile_pipeline_dsl_artifact_with_controllers(dsl, controller_manifests()).graph.to_dict()
return run_cv_refit_bundle(
dsl=dsl, envelope=envelope, graph=graph, dataset_path=dataset_arg, workdir=run_dir, dagml_cli=cli, venv_python=venv_python, selection_metric=metric, dataset_pickle=dataset_pickle, dataset=spectro, random_state=random_state, refit=refit, refit_top_k=refit_top_k
)
outcome = run_cv_refit_lane(
steps=steps, dsl=dsl, envelope=envelope, spectro=spectro, identity=identity, excluded=excluded,
selection_metric=metric, refit=refit, refit_top_k=refit_top_k, run_callback=run_callback,
)
if outcome["returncode"] != 0:
_raise_run_failure(outcome, "dag-ml engine run failed")
# A native param sweep varies one model's params only (same model class across variants), so only the
# per-variant config_name map is needed — model_name is shared, so the projection's scalar fallback is
# correct. The map keys each opaque native variant_id to its legacy expand config name. A NON-degenerate
# sweep (`_grid_`) selects the true CV-best, whose config_name is the WINNING variant's name (NOT
# index 0): recover it by matching the winner's refit model params against the per-variant model params
# (aligned with `variant_config_names`), so the winner is content-paired exactly like the operator path.
catalog_map = _native_param_config_map_from_catalog(
outcome.get("variant_catalog", []), variant_config_names, variant_model_params,
)
if not refit and variant_config_names and len(catalog_map) != len(variant_config_names):
raise DagMlUnsupported("native CV-only parameter sweep did not expose an unambiguous variant-to-parameter mapping")
winner_config_name = _native_param_winner_config_name(outcome["refit_artifacts"], variant_config_names, variant_model_params)
variant_config_map = (
catalog_map or _native_variant_config_map(outcome["scores"], variant_config_names, winner_config_name)
if variant_config_names else None
)
# FILL the strict direct-block rows with this run's per-sample y_pred/y_true/sample_indices (the
# winner's refit `(final, train)` + `(final, test)` + per-fold OOF; each loser's per-fold OOF). dag-ml
# emits ONE bundle here with every variant's reports correctly tagged (only the winner refits), so the
# NodeResult frames key per-variant exactly like the reports — bucket them with `_frames_by_variant`
# under each frame's own `variant_id` (the winner's untagged OOF-avg frame falls back to the winner id,
# the `(final, None)`-report owner). The projection then reads each row's arrays from ITS OWN variant's
# blocks (no cross-variant y_pred leakage), exactly as the operator-sweep / single-pipeline paths do.
winner_variant_id = next(
(report.get("variant_id") for report in (outcome["scores"] or {}).get("reports", []) if report["partition"] == "final" and report.get("fold_id") is None),
None,
)
if winner_variant_id is None and not refit:
winner_variant_id = next(
(report.get("variant_id") for report in (outcome["scores"] or {}).get("reports", []) if report["partition"] == "validation" and report.get("fold_id") != "avg"),
None,
)
results_by_variant = _frames_by_variant(outcome["results"], winner_variant_id)
result = _scores_to_run_result(
outcome["scores"], spectro.name, _model_name(steps), metric, task_type, config_name=config_name, variant_config_names=variant_config_map, results_by_variant=results_by_variant, identity=identity, refit_artifacts=outcome["refit_artifacts"], emit_all_refits=refit_top_k > 1, refit_name_suffix=f"_refit_rmsecvt{refit_top_k}" if refit_top_k > 1 else "_refit"
)
if refit_top_k > 1:
result.per_dataset[spectro.name]["selected_refit_variant_ids"] = outcome.get("selected_refit_variant_ids", [])
record_execution_lane(result, outcome["execution_lane"])
return result
def _run_native_operator_generation(
pipeline: list[Any],
spectro: Any,
dataset_arg: str,
cli: str,
venv_python: str,
run_dir: Path,
metric: str,
task_type: str,
cv_pool: list[int] | None = None,
excluded: set[int] | None = None,
tags_by_sample: dict[int, list[str]] | None = None,
dataset_pickle: str | None = None,
config_name: str = "",
variant_config_names: list[str] | None = None,
random_state: int | None = None,
refit: bool = True,
refit_top_k: int = 1,
) -> RunResult:
"""Run a FLAT-SINGLE operator ``_or_`` as ONE native dag-ml operator-SELECT + refit run (#23 Phase 7).
The generator sits on a TRANSFORM step (the model is concrete): the bridge lowers the ``_or_`` to a
compat ``Generator`` step, dag-ml's ``compile_operator_variant_models`` expands the operator-variant
models, and the in-process binding scores EACH choice by its cross-fold OOF ``metric``, refits the
requested top-k candidates on their own pruned plans, and surfaces every variant's validation reports.
Each report carries the cross-language ``variant_label`` content fingerprint (the WINNER too).
``bundle.scores`` is mapped to the full
PER-VARIANT legacy table, keyed CONTENT-WISE (``variant_label`` → ``config_name``), so a sweep's
num_predictions + winner identity match the Python-expand path.
Differs from :func:`_run_native_generation` (the param-sweep template) in three ways: (1) NO
``_apply_plain_model_params`` — the generator is on the transform, the model is already concrete; (2)
the union graph has ONE model node PER choice (Mechanism B namespaces them), so EVERY model node needs
a data binding (mirroring the separation-branch fan-out), not just the first; (3) the variant-config
map is keyed by ``variant_label``
(:func:`~nirs4all.pipeline.dagml.result._native_operator_variant_config_map`), not the positional zip
a param sweep uses.
FALLBACK CONTRACT (the inner Python-expand fallback fires on LOWERING-UNSUPPORTED ONLY): every
lowering step — splitter check, the bridge ``_or_`` lowering (:func:`assemble_cv_refit_dsl`), and the
per-choice ``variant_label`` fingerprinting (:func:`_native_operator_config_by_label`) — runs inside a
narrow guard that converts a lowering refusal (``NotImplementedError`` / :class:`DagMlUnsupported`) into
a DISTINCT :class:`~.errors._OperatorLoweringUnsupported` sentinel, which the routing branch catches →
Python-expand. Everything AFTER the guard (compile / run / result-mapping) is OUTSIDE it, so a genuine
runtime error there PROPAGATES — including a runtime ``DagMlUnsupported`` from :func:`_raise_run_failure`
(a non-zero run classified ``error_kind == "unsupported"``), which is a REAL coverage boundary the host
did not pre-check, NOT a lowering gap. The routing branch catches ONLY the sentinel, so that runtime
``DagMlUnsupported`` is never silently reclassified as lowering-unsupported and masked.
"""
import dag_ml
from .cli_runner import data_bindings_for_nodes
from .errors import _OperatorLoweringUnsupported
from .result import _frames_by_variant, _native_operator_config_by_label, _native_operator_variant_config_map
# --- LOWERING PHASE (narrow fallback scope) ---------------------------------------------------------
# Only a LOWERING refusal demotes to Python-expand. The bridge `_or_` lowering raises NotImplementedError;
# the label fingerprinting raises DagMlUnsupported; both are lowering-unsupported, so convert them to the
# DISTINCT `_OperatorLoweringUnsupported` sentinel the routing branch catches. A non-lowering error (e.g.
# an internal invariant) is NOT a coverage gap — let it propagate.
try:
from nirs4all.pipeline.config._generator.keywords import GENERATION_KEYWORDS
steps, splitter = _split_pipeline(pipeline)
if splitter is None:
raise DagMlUnsupported("engine='dag-ml' requires a cross-validator step (e.g. KFold) in the pipeline")
_reject_multi_model(steps)
# Reject UP FRONT the transform-side operators the X-chain cannot run — but skip the generator step
# itself (its bare-operator choices were already routability-gated by the predicate, and are
# fingerprinted below). The concrete model is validated by its own fit, like every other path.
_assert_supported_operators([step for step in steps if not (isinstance(step, dict) and GENERATION_KEYWORDS & set(step))])
identity = mint_identity(spectro)
pool = list(cv_pool) if cv_pool is not None else spectro.index_column("sample", {"partition": "train"})
folds = _build_folds(splitter, spectro, pool, excluded or set())
envelope = build_envelope(spectro, identity, sample_ints=pool, excluded_sample_ints=excluded or None, tags_by_sample=tags_by_sample, group_by_sample=_split_group_grain(splitter, spectro, pool))
# The DSL carries the lowered Generator on the transform position. A FLAT-SINGLE bare `_or_` lowers
# via the compat fusion (assemble_cv_refit_dsl → pipeline_to_dsl); a CONSTRAINED `_or_`-pick /
# `_cartesian_` OR an UNCONSTRAINED pick/arrange/`_cartesian_` (ADR-17 item 5 slice C) lowers each
# survivor into ONE model-terminated canonical Generator branch (assemble_constrained_cv_refit_dsl):
# both produce multi-op survivor SEQUENCES the compat fusion's operator-variant compiler refuses for a
# model-free choice, and the native generator carries the pick/arrange selectors (+ constraints when
# present) so dag-ml prunes the byte-identical set. Both raise from the bridge on an unlowerable shape.
# Compute the content-keyed {variant_label -> config_name} map HERE too (it fingerprints each survivor
# — a lowering step that can raise DagMlUnsupported on a non-finite / non-JSON param), so a label
# failure demotes BEFORE the run.
if _is_constrained_operator_generator(pipeline) or _is_unconstrained_operator_generator(pipeline):
dsl = assemble_constrained_cv_refit_dsl(steps, identity, envelope, folds, dsl_id="nirs4all-pipeline", n_splits=len(folds))
else:
dsl = assemble_cv_refit_dsl(steps, identity, envelope, folds, dsl_id="nirs4all-pipeline", n_splits=len(folds))
config_by_label = _native_operator_config_by_label(steps, variant_config_names or [])
except (DagMlUnsupported, NotImplementedError) as exc:
raise _OperatorLoweringUnsupported(f"operator generator lowering unsupported, demoting to Python expand: {exc}") from exc
# --- COMPILE / RUN / RESULT-MAPPING (errors PROPAGATE — never reclassified as a coverage gap) --------
# Operator SELECT derives seeded variants; bind that seed into the campaign
# before compilation so pruning preserves their strict native identities.
dsl["root_seed"] = random_state if random_state is not None else 0
# The union graph compiles ONE model node per `_or_` choice (Mechanism B namespacing), so bind EVERY
# model node — a single binding on the first node would leave the other choices' model nodes with empty
# data_views (the separation-branch fan-out path solves the same multi-model-node binding the same way).
graph = dag_ml.compile_pipeline_dsl_artifact_with_controllers(dsl, controller_manifests()).graph.to_dict()
model_ids = [node["id"] for node in graph["nodes"] if node["kind"] == "model"]
dsl["data_bindings"] = data_bindings_for_nodes(model_ids, envelope)
outcome = run_cv_refit_bundle(
dsl=dsl, envelope=envelope, graph=graph, dataset_path=dataset_arg, workdir=run_dir, dagml_cli=cli, venv_python=venv_python, selection_metric=metric, dataset_pickle=dataset_pickle, dataset=spectro, random_state=random_state, refit=refit, refit_top_k=refit_top_k
)
if outcome["returncode"] != 0:
_raise_run_failure(outcome, "dag-ml operator-generation run failed")
# CONTENT-keyed map: every report carries its choice's `variant_label`, so each variant's `variant_id`
# resolves to its legacy `config_name` by the pre-computed fingerprint map (NOT a positional zip). The
# winner's reports (val + refit) are threaded under its OWN variant_id so its per-fold + final/test rows
# carry real per-sample y_pred (2a-i/ii).
scores = outcome["scores"]
variant_config_map = _native_operator_variant_config_map(scores, config_by_label)
winner_variant_id = next(
(report.get("variant_id") for report in (scores or {}).get("reports", []) if report["partition"] == "final" and report.get("fold_id") is None),
None,
)
if winner_variant_id is None and not refit:
winner_variant_id = next(
(report.get("variant_id") for report in (scores or {}).get("reports", []) if report["partition"] == "validation" and report.get("fold_id") != "avg"),
None,
)
# Split the surfaced frames PER VARIANT so a LOSER variant's per-fold val rows fill from ITS OWN
# validation (OOF) predictions, not just the winner's. dag-ml surfaces each loser's per-fold val
# blocks re-tagged with the loser's variant_id (top-level in-process / `lineage.variant_id`
# subprocess); the winner's real per-fold + refit frames carry the winner's id. Untagged frames
# (the winner's OOF-average frame) default to the winner. NO cross-variant leakage: a frame routes
# to its OWN variant only.
results_by_variant = _frames_by_variant(outcome["results"], winner_variant_id) if winner_variant_id is not None else None
result = _scores_to_run_result(
scores, spectro.name, _model_name(steps), metric, task_type, config_name=config_name, variant_config_names=variant_config_map or None, results_by_variant=results_by_variant, identity=identity, refit_artifacts=outcome["refit_artifacts"], emit_all_refits=refit_top_k > 1, refit_name_suffix=f"_refit_rmsecvt{refit_top_k}" if refit_top_k > 1 else "_refit"
)
if refit_top_k > 1:
result.per_dataset[spectro.name]["selected_refit_variant_ids"] = outcome.get("selected_refit_variant_ids", [])
return result
def _bind_experimental_unit_design(dsl: dict[str, Any], spectro: Any, identity: Any, pool: list[int]) -> None:
"""Sign explicit units before compilation; native tasks own scope weights."""
if getattr(spectro, "independent_unit_ids", None) is None:
return
from .envelope import target_names
from .experimental_units import apply_experimental_unit_contract
from .resolver import MaterializationResolver
targets = MaterializationResolver(spectro, identity).resolve_targets([identity.to_wire(sample) for sample in pool])
names = targets.get("target_names", ["y"] if len(target_names(spectro)) == 1 else target_names(spectro))
apply_experimental_unit_contract(dsl, spectro, identity, pool, target_values=targets["values"], target_names=names)
def _run_concrete_scores(
pipeline: Any,
spectro: Any,
dataset_arg: str,
cli: str,
venv_python: str,
run_dir: Path,
cv_pool: list[int] | None = None,
excluded: set[int] | None = None,
tags_by_sample: dict[int, list[str]] | None = None,
dataset_pickle: str | None = None,
random_state: int | None = None,
refit: bool = True,
metric: str = "rmse",
) -> tuple[dict[str, Any], str, list[dict[str, Any]], Any, list[dict[str, Any]], dict[str, Any]]:
"""Run one concrete (generator-free) pipeline; return ``(scores, model_name, results, identity, refit_artifacts, lane)``.
Return native scores, model label, per-node results, sample identities,
captured fitted estimators and the execution-lane record (see
:mod:`.methods_lane`). The projection preserves all native measurements,
independent of historical legacy serialization/refit behavior. ``cv_pool``
is the CV sample universe; ``excluded`` is marked in the envelope only when
``keep_in_oof=True``.
"""
steps, splitter = _split_pipeline(pipeline)
if splitter is None:
raise DagMlUnsupported("engine='dag-ml' requires a cross-validator step (e.g. KFold) in the pipeline")
_reject_multi_model(steps)
_assert_supported_operators(steps)
steps = _apply_model_params(steps)
identity = mint_identity(spectro)
pool = list(cv_pool) if cv_pool is not None else spectro.index_column("sample", {"partition": "train"})
folds = _build_folds(splitter, spectro, pool, excluded or set())
envelope = build_envelope(spectro, identity, sample_ints=pool, excluded_sample_ints=excluded or None, tags_by_sample=tags_by_sample, group_by_sample=_split_group_grain(splitter, spectro, pool))
dsl = assemble_cv_refit_dsl(
steps, identity, envelope, folds, dsl_id="nirs4all-pipeline", n_splits=len(folds),
force_fitted_x_chain=getattr(spectro, "_generated_view_store", None) is not None and len(steps) == 2,
)
def run_callback() -> dict[str, Any]:
import dag_ml
from nirs4all.pipeline.dagml_bridge import named_model_input_spec
_bind_experimental_unit_design(dsl, spectro, identity, pool)
manifests = controller_manifests(dsl)
dsl["root_seed"] = random_state if random_state is not None else 0
graph = dag_ml.compile_pipeline_dsl_artifact_with_controllers(dsl, manifests).graph.to_dict()
return run_cv_refit_bundle(
dsl=dsl, envelope=envelope, graph=graph, dataset_path=dataset_arg, workdir=run_dir, dagml_cli=cli, venv_python=venv_python, selection_metric=metric, dataset_pickle=dataset_pickle, dataset=spectro, random_state=random_state, refit=refit
)
outcome = run_cv_refit_lane(
steps=steps, dsl=dsl, envelope=envelope, spectro=spectro, identity=identity, excluded=excluded,
selection_metric=metric, refit=refit, refit_top_k=1, run_callback=run_callback,
)
if outcome["returncode"] != 0:
_raise_run_failure(outcome, "dag-ml engine run failed")
return outcome["scores"], _model_name(steps), outcome["results"], identity, outcome["refit_artifacts"], outcome["execution_lane"]
def _run_concrete(
pipeline: Any,
spectro: Any,
dataset_arg: str,
cli: str,
venv_python: str,
run_dir: Path,
metric: str = "rmse",
task_type: str = "regression",
cv_pool: list[int] | None = None,
excluded: set[int] | None = None,
tags_by_sample: dict[int, list[str]] | None = None,
dataset_pickle: str | None = None,
config_name: str = "",
random_state: int | None = None,
refit: bool = True,
) -> RunResult:
"""Run one concrete (generator-free) pipeline through dag-ml-cli; map its native scores.
``cv_pool`` is the CV sample-int universe (de-excluded pool in legacy mode, full train in opt-in
mode); ``excluded`` is marked in the envelope only in the opt-in (``keep_in_oof=True``) mode.
"""
scores, model_name, results, identity, refit_artifacts, lane = _run_concrete_scores(
pipeline, spectro, dataset_arg, cli, venv_python, run_dir, cv_pool, excluded, tags_by_sample, dataset_pickle=dataset_pickle, random_state=random_state, refit=refit, metric=metric
)
result = _scores_to_run_result(scores, spectro.name, model_name, metric, task_type, config_name=config_name, results=results, identity=identity, refit_artifacts=refit_artifacts)
record_execution_lane(result, lane)
return result
def _source_concat_layout(source_indices: list[int], spectro: Any, envelope: dict[str, Any] | None = None) -> dict[str, Any]:
"""Internal source-layout marker for a proven ``{"merge": {"sources": "concat"}}`` boundary."""
envelope_layout = ((envelope or {}).get("plan") or {}).get("source_layout")
if isinstance(envelope_layout, dict) and source_indices == list(range(spectro.features_sources())):
return dict(envelope_layout)
source_order = [f"src{index}" for index in source_indices]
blocks: list[dict[str, Any]] = []
start = 0
num_features = spectro.num_features
for source_id, source_index in zip(source_order, source_indices, strict=True):
width = int(num_features[source_index] if isinstance(num_features, list) else num_features)
blocks.append(
{
"source_id": source_id,
"source_index": source_index,
"column_start": start,
"column_count": width,
}
)
start += width
return {
"kind": "by_source_concat",
"source_order": source_order,
"source_indices": list(source_indices),
"blocks": blocks,
"concat": {
"axis": "feature",
"total_column_count": start,
"preserve_source_order": True,
},
}
def _graph_upstream_x_chain(graph: dict[str, Any], node_id: str) -> list[str]:
"""Ordered upstream x-chain ids for a compiled linear graph."""
incoming: dict[str, str] = {}
for edge in graph.get("edges", []) or []:
target = edge.get("target") or {}
source = edge.get("source") or {}
contract = edge.get("contract") or {}
if target.get("port_name") == "x" and contract.get("kind") == "data":
incoming[str(target["node_id"])] = str(source["node_id"])
chain: list[str] = []
current = incoming.get(node_id)
while current is not None:
chain.append(current)
current = incoming.get(current)
chain.reverse()
return chain
def _source_concat_preprocessing_metadata(
graph: dict[str, Any],
model_node: dict[str, Any],
source_indices: list[int],
source_layout: dict[str, Any],
) -> dict[str, Any]:
"""Build the node-runner contract for top-level source concat preprocessing."""
nodes_by_id = {node["id"]: node for node in graph.get("nodes", []) or []}
chain_nodes = [nodes_by_id[node_id] for node_id in _graph_upstream_x_chain(graph, str(model_node["id"]))]
if any(node.get("kind") != "transform" for node in chain_nodes):
raise DagMlUnsupported("source-concat preprocessing chain may contain only X transform nodes")
blocks_by_index = {
int(block["source_index"]): block
for block in source_layout.get("blocks", []) or []
if isinstance(block, dict) and "source_index" in block
}
raw_source_ids = source_layout.get("source_ids")
raw_source_order = source_layout.get("source_order")
source_ids: list[Any] = raw_source_ids if isinstance(raw_source_ids, list) else []
source_order: list[Any] = raw_source_order if isinstance(raw_source_order, list) else []
sources: list[dict[str, Any]] = []
for position, source_index in enumerate(source_indices):
block = blocks_by_index.get(source_index, {})
sources.append(
{
"source_name": block.get("source_name", source_order[position] if position < len(source_order) else f"source_{source_index}"),
"source_id": block.get("source_id", source_ids[position] if position < len(source_ids) else f"src{source_index}"),
"source_index": source_index,
"steps": chain_nodes,
}
)
return {
"mode": "top_level_sources_concat",
"preserve_legacy_sources_after_merge": False,
"source_layout": source_layout,
"sources": sources,
}
def _mark_source_concat_model_nodes(graph: dict[str, Any], source_indices: list[int], spectro: Any, envelope: dict[str, Any]) -> None:
"""Attach source-concat layout and per-source preprocessing metadata to model nodes."""
layout = _source_concat_layout(source_indices, spectro, envelope)
for node in graph.get("nodes", []):
if node.get("kind") != "model":
continue
metadata = dict(node.get("metadata") or {})
metadata["source_layout"] = layout
metadata["source_concat_preprocessing"] = _source_concat_preprocessing_metadata(graph, node, source_indices, layout)
node["metadata"] = metadata
def _run_source_concat_merge(
pre_merge_steps: list[Any],
post_merge_steps: list[Any],
source_indices: list[int],
spectro: Any,
dataset_arg: str,
cli: str,
venv_python: str,
run_dir: Path,
metric: str,
task_type: str,
dataset_pickle: str | None = None,
config_name: str = "",
random_state: int | None = None,
) -> RunResult:
"""Run top-level source concat natively by preserving the per-source transform boundary.
Legacy applies upstream X transforms to each source independently, then ``{"merge": {"sources":
"concat"}}`` hstacks those transformed blocks before the downstream splitter/model. A plain native
early-fusion run would fit the X-chain on the already-concatenated matrix, which changes row-wise
transforms such as SNV. This path compiles the graph without the merge step but marks the model node
with a source-layout contract; the node runner then materializes source blocks, applies the upstream
X-chain per block, hstacks in source order, and fits/predicts the model.
"""
native_pipeline = [*pre_merge_steps, *post_merge_steps]
steps, splitter = _split_pipeline(native_pipeline)
if splitter is None:
raise DagMlUnsupported("engine='dag-ml' requires a cross-validator step (e.g. KFold) in the pipeline")
_reject_multi_model(steps)
_assert_supported_operators(steps)
steps = _apply_model_params(steps)
identity = mint_identity(spectro)
pool = spectro.index_column("sample", {"partition": "train"})
folds = _build_folds(splitter, spectro, pool, set())
envelope = build_envelope(spectro, identity, sample_ints=pool, group_by_sample=_split_group_grain(splitter, spectro, pool))
dsl = assemble_cv_refit_dsl(steps, identity, envelope, folds, dsl_id="nirs4all-source-concat", n_splits=len(folds))
import dag_ml
dsl["root_seed"] = random_state if random_state is not None else 0
graph = dag_ml.compile_pipeline_dsl_artifact_with_controllers(dsl, controller_manifests()).graph.to_dict()
_mark_source_concat_model_nodes(graph, source_indices, spectro, envelope)
outcome = run_cv_refit_bundle(
dsl=dsl,
envelope=envelope,
graph=graph,
dataset_path=dataset_arg,
workdir=run_dir,
dagml_cli=cli,
venv_python=venv_python,
selection_metric=metric,
dataset_pickle=dataset_pickle,
dataset=spectro,
random_state=random_state,
)
if outcome["returncode"] != 0:
_raise_run_failure(outcome, "dag-ml source-concat merge run failed")
return _scores_to_run_result(
outcome["scores"],
spectro.name,
_model_name(steps),
metric,
task_type,
config_name=config_name,
results=outcome["results"],
identity=identity,
refit_artifacts=outcome["refit_artifacts"],
)
def _run_repetition(pipeline: list[Any], spectro: Any, dataset_arg: str, cli: str, venv_python: str, run_dir: Path, metric: str, task_type: str, dataset_pickle: str | None = None, config_name: str = "", random_state: int | None = None, refit: bool = True) -> RunResult:
"""Run a REPETITION (sample-grain grouped) pipeline as ONE native dag-ml CV+refit run.
The CV universe is the repetition ROWS of the train partition (each stored row is its own
sample int / sample_id — repetitions are NOT collapsed). Folds are GROUP-aware
(:func:`_build_group_folds`): all replicates of a sample land on the same fold side, so every
rep row is validated exactly once (a clean OOF partition) while a group is never split. The
envelope carries each row's ``group_id`` (the repetition column value), and dag-ml-data's
``validate_fold_set_against_sample_relations`` refuses any fold that splits a group — the
native group-leakage guarantee. Each rep row is scored individually (the repetition grain),
which is exactly what nirs4all's ``cv_best_score``/``best_rmse`` report; the sample-level
aggregation (the ``_agg`` twin) is a separate score nirs4all does not surface on RunResult,
so no aggregation reducer is needed.
Generators are expanded in Python (operator-level via ``expand_spec``; a param-level model
sweep also goes through ``expand_spec`` here for simplicity) and each concrete variant runs
through the group-aware path, selecting the best by its CV score — mirroring nirs4all.
"""
from nirs4all.pipeline.config.generator import expand_spec
variants = expand_spec(pipeline)
results = [
_run_repetition_concrete(variant, spectro, dataset_arg, cli, venv_python, run_dir / f"variant{index}", metric, task_type, dataset_pickle=dataset_pickle, config_name=config_name, random_state=random_state, refit=refit)
for index, variant in enumerate(variants)
]
if len(results) == 1:
return results[0]
# Direction from the SINGLE source of truth (core.metrics), so a classification sweep ranking on
# `balanced_accuracy` (the default since #60) is MAXIMIZED, not minimized (which would pick the worst
# variant). The old `metric in ("accuracy", "r2")` set silently mis-ranked balanced_accuracy.
maximize = is_higher_better(metric)
def _cv_rank(result: RunResult) -> float:
score = result.cv_best_score
if score != score: # NaN ranks last
return float("inf")
return -score if maximize else score
return min(results, key=_cv_rank)
def _run_repetition_concrete(pipeline: Any, spectro: Any, dataset_arg: str, cli: str, venv_python: str, run_dir: Path, metric: str, task_type: str, dataset_pickle: str | None = None, config_name: str = "", random_state: int | None = None, refit: bool = True) -> RunResult:
"""One concrete repetition variant: group-aware folds + a ``group_id``-carrying envelope."""
from .exclude import _resolve_exclude
steps, splitter = _split_pipeline(pipeline)
if splitter is None:
raise DagMlUnsupported("engine='dag-ml' requires a cross-validator step (e.g. KFold) in the pipeline")
steps, pool, excluded = _resolve_exclude(steps, spectro)
_assert_supported_operators(steps)
steps = _apply_model_params(steps)
identity = mint_identity(spectro)
folds = _build_group_folds(splitter, spectro, pool)
if isinstance(excluded, FoldLocalExclusion):
folds = excluded.apply(spectro, folds)
group_by_sample = _split_group_grain(splitter, spectro, pool) or _repetition_grain(spectro, pool)
envelope = build_envelope(spectro, identity, sample_ints=pool, excluded_sample_ints=excluded, group_by_sample=group_by_sample)
dsl = assemble_cv_refit_dsl(steps, identity, envelope, folds, dsl_id="nirs4all-pipeline", n_splits=len(folds))
import dag_ml
dsl["root_seed"] = random_state if random_state is not None else 0
graph = dag_ml.compile_pipeline_dsl_artifact_with_controllers(dsl, controller_manifests()).graph.to_dict()
outcome = run_cv_refit_bundle(
dsl=dsl, envelope=envelope, graph=graph, dataset_path=dataset_arg, workdir=run_dir, dagml_cli=cli, venv_python=venv_python, selection_metric=metric, dataset_pickle=dataset_pickle, dataset=spectro, random_state=random_state, refit=refit
)
if outcome["returncode"] != 0:
_raise_run_failure(outcome, "dag-ml repetition run failed")
# Thread the node results + minted identity into the projection so the strict direct-block rows
# (per-fold val + refit final/test) carry real y_pred/y_true/sample_indices — the same 2a-i fill the
# single-pipeline path does. For a SAMPLE-LEVEL-aggregation dataset (`aggregate=True`) the refit's
# `(test, None)` block dag-ml emits is already the aggregated sample grain (the replicates were
# collapsed at materialization), so this surfaces the aggregation's final-(test) y_pred at parity with
# legacy (Gap 2). Native measurements determine the projected rows.
result = _scores_to_run_result(outcome["scores"], spectro.name, _model_name(steps), metric, task_type, config_name=config_name, results=outcome["results"], identity=identity, refit_artifacts=outcome["refit_artifacts"])
from .native_vote import project_vote_evidence
project_vote_evidence(result, outcome.get("classification_evidence", []), identity)
_attach_repetition_prediction_context(result, spectro)
_add_explicit_aggregate_twins(result, spectro)
return result
def _attach_repetition_prediction_context(result: RunResult, spectro: Any) -> None:
"""Carry dataset repetition/aggregation defaults onto the projected prediction table."""
if getattr(spectro, "repetition", None):
result.predictions.set_repetition_column(spectro.repetition)
result.predictions.set_aggregate_context(
getattr(spectro, "aggregate", None),
getattr(spectro, "aggregate_method", None),
bool(getattr(spectro, "aggregate_exclude_outliers", False)),
)
def _add_explicit_aggregate_twins(result: RunResult, spectro: Any) -> None:
"""Append legacy-style ``*_agg`` companion rows for explicit repetition aggregates.
The legacy orchestrator persists display-only aggregate twins only when the dataset names a concrete
aggregate metadata column. Native dag-ml keeps the ranking surface on raw rows, but the parity registry
still expects the public row inventory to include those companion rows.
"""
aggregate = getattr(spectro, "aggregate", None)
if not aggregate or aggregate == "y":
return
method = getattr(spectro, "aggregate_method", None) or "mean"
exclude_outliers = bool(getattr(spectro, "aggregate_exclude_outliers", False))
source_entries = list(result.predictions.iter_entries())
for entry in source_entries:
fold_id = str(entry.get("fold_id") or "")
if not fold_id or fold_id.endswith("_agg"):
continue
_add_explicit_aggregate_twin(result.predictions, spectro, entry, str(aggregate), method, exclude_outliers)
def _add_explicit_aggregate_twin(predictions: Predictions, spectro: Any, entry: dict[str, Any], aggregate: str, method: str, exclude_outliers: bool) -> None:
fold_id = str(entry.get("fold_id") or "")
partition = entry.get("partition", "")
primary_metric = entry.get("metric") or "rmse"
metric_names = _aggregate_metric_names(entry, primary_metric)
payload = _aggregate_entry_payload(spectro, entry, aggregate, method, exclude_outliers, metric_names)
if payload is None:
payload = _score_only_aggregate_payload(entry, partition, primary_metric)
agg_y_true, agg_y_pred, agg_y_proba, agg_scores, primary_value, n_samples = payload
result_metadata = copy.deepcopy(entry.get("result_metadata") or {})
projection = result_metadata.get("dagml_projection")
if isinstance(projection, dict):
score_provenance = projection.setdefault("score_provenance", {})
source = copy.deepcopy(score_provenance.get(partition) or {})
source.update({
"purpose": "measurement",
"derived_by": "nirs4all.data.Predictions.aggregate",
"source_prediction_id": entry.get("id"),
"aggregate_column": aggregate,
"aggregate_method": method,
})
score_provenance[partition] = source
result_metadata["aggregate_evidence"] = {
"owner": "nirs4all.data.Predictions.aggregate", "source_prediction_id": entry.get("id"),
"column": aggregate, "method": method, "partition": partition, "selection_score": False,
}
predictions.add_prediction(
dataset_name=entry.get("dataset_name", ""),
dataset_path=entry.get("dataset_path", ""),
config_name=entry.get("config_name", ""),
config_path=entry.get("config_path", ""),
pipeline_uid=entry.get("pipeline_uid"),
step_idx=entry.get("step_idx", 0),
op_counter=entry.get("op_counter", 0),
model_name=entry.get("model_name", ""),
model_classname=entry.get("model_classname", ""),
model_path=entry.get("model_path", ""),
fold_id=f"{fold_id}_agg",
sample_indices=None,
weights=None,
metadata={},
result_metadata=result_metadata,
partition=partition,
y_true=agg_y_true,
y_pred=agg_y_pred,
y_proba=agg_y_proba,
val_score=primary_value if partition in ("val", "validation") else entry.get("val_score"),
test_score=primary_value if partition == "test" else entry.get("test_score"),
train_score=primary_value if partition == "train" else entry.get("train_score"),
metric=primary_metric,
task_type=entry.get("task_type", "regression"),
n_samples=n_samples,
n_features=entry.get("n_features", 0),
preprocessings=entry.get("preprocessings", ""),
best_params=entry.get("best_params") or {},
scores={partition: agg_scores} if partition else {"": agg_scores},
branch_id=entry.get("branch_id"),
branch_path=entry.get("branch_path"),
branch_name=entry.get("branch_name"),
exclusion_count=entry.get("exclusion_count"),
exclusion_rate=entry.get("exclusion_rate"),
model_artifact_id=entry.get("model_artifact_id"),
trace_id=entry.get("trace_id"),
refit_context=entry.get("refit_context"),
target_processing=entry.get("target_processing", ""),
)
def _aggregate_metric_names(entry: dict[str, Any], primary_metric: str) -> list[str]:
for part_scores in (entry.get("scores") or {}).values():
if isinstance(part_scores, dict) and part_scores:
metric_names = [str(metric_name) for metric_name in part_scores if ":" not in str(metric_name)]
return metric_names or [primary_metric]
return [primary_metric]
def _score_only_aggregate_payload(entry: dict[str, Any], partition: str, primary_metric: str) -> tuple[np.ndarray, np.ndarray, None, dict[str, float], float | None, int]:
partition_scores = (entry.get("scores") or {}).get(partition)
if not isinstance(partition_scores, dict):
partition_scores = {}
scores = {str(metric): float(value) for metric, value in partition_scores.items() if value is not None}
return np.array([]), np.array([]), None, scores, scores.get(primary_metric), 0
def _aggregate_entry_payload(
spectro: Any,
entry: dict[str, Any],
aggregate: str,
method: str,
exclude_outliers: bool,
metric_names: list[str],
) -> tuple[np.ndarray, np.ndarray, np.ndarray | None, dict[str, float], float | None, int] | None:
if entry.get("y_true") is None or entry.get("y_pred") is None:
return None
y_true = np.asarray(entry.get("y_true")).flatten()
y_pred = np.asarray(entry.get("y_pred")).flatten()
if y_true.size == 0 or y_pred.size == 0:
return None
group_ids = _aggregate_group_ids(spectro, entry, aggregate, y_pred.size)
if group_ids is None or len(np.unique(group_ids)) == y_pred.size:
return None
y_proba_raw = entry.get("y_proba")
y_proba = None
if y_proba_raw is not None:
y_proba_candidate = np.asarray(y_proba_raw)
if y_proba_candidate.size > 0:
y_proba = y_proba_candidate
try:
aggregate_result = Predictions.aggregate(
y_pred=y_pred,
group_ids=group_ids,
y_true=y_true,
y_proba=y_proba,
method=method,
exclude_outliers=exclude_outliers,
)
except Exception:
return None
agg_y_true = aggregate_result.get("y_true")
agg_y_pred = aggregate_result.get("y_pred")
if agg_y_true is None or agg_y_pred is None or len(agg_y_true) == 0:
return None
classes = (entry.get("result_metadata") or {}).get("classification_evidence", {}).get("classes")
if y_proba is not None and classes is not None:
# The aggregation owner returns probability-column positions. Their
# explicit classifier labels, not positional integers, are the targets.
agg_y_pred = np.asarray(classes)[np.asarray(agg_y_pred, dtype=int)]
try:
metric_values = eval_list(agg_y_true, agg_y_pred, metric_names)
except Exception:
return None
scores = {
metric_name: float(value)
for metric_name, value in zip(metric_names, metric_values, strict=False)
if value is not None
}
primary_metric = entry.get("metric") or "rmse"
return (
np.asarray(agg_y_true),
np.asarray(agg_y_pred),
aggregate_result.get("y_proba"),
scores,
scores.get(primary_metric),
int(len(agg_y_true)),
)
def _aggregate_group_ids(spectro: Any, entry: dict[str, Any], aggregate: str, expected_size: int) -> np.ndarray | None:
sample_indices = entry.get("sample_indices") or []
if len(sample_indices) != expected_size:
return None
sample_ints = [int(sample_int) for sample_int in sample_indices]
try:
values = np.asarray(spectro.metadata_column(aggregate, {"sample": sample_ints}, include_augmented=False))
stored = [int(sample_int) for sample_int in spectro.index_column("sample", {"sample": sample_ints})]
except Exception:
return None
if len(values) != len(stored):
return None
row_of = {sample_int: row for row, sample_int in enumerate(stored)}
try:
return np.asarray([values[row_of[sample_int]] for sample_int in sample_ints])
except KeyError:
return None
def _reshape_for_rep_fusion(rep_step: dict[str, Any], spectro: Any) -> None:
"""Run nirs4all's REAL repetition reshape (``rep_to_sources`` / ``rep_to_pp``) in place.
Reuses the production operator (:class:`RepetitionConfig`) + the dataset's own reshape methods
(``reshape_reps_to_sources`` / ``reshape_reps_to_preprocessings``) — no reshape logic is
reimplemented here. The reshape collapses the N replicate ROWS of a physical sample into either
N sample-aligned feature SOURCES (``rep_to_sources``) or N stacked PROCESSING layers
(``rep_to_pp``), reducing the dataset to ``n_unique`` physical samples.
The ``repetition`` flag is cleared afterwards: it named the rep grouping of the ORIGINAL rows,
which no longer exists once the replicates have become sources/processings — the reshaped
dataset's unit of analysis is the physical sample, so the downstream folds must be sample-grain
(the plain-repetition group-fold path, #21, must NOT fire on the reshaped dataset).
"""
from nirs4all.operators.data.repetition import RepetitionConfig
if "rep_to_sources" in rep_step:
config = RepetitionConfig.from_step_value(rep_step["rep_to_sources"])
spectro.reshape_reps_to_sources(config)
else:
config = RepetitionConfig.from_step_value(rep_step["rep_to_pp"])
spectro.reshape_reps_to_preprocessings(config)
spectro._repetition = None # noqa: SLF001 - the rep grouping was consumed by the reshape
def _run_rep_fusion(
pipeline: list[Any],
rep_step: dict[str, Any],
spectro: Any,
dataset_arg: str,
cli: str,
venv_python: str,
run_dir: Path,
metric: str,
task_type: str,
config_name: str = "",
variant_config_names: list[str] | None = None,
is_classification: bool = False,
random_state: int | None = None,
) -> RunResult:
"""Run a ``rep_to_sources`` / ``rep_to_pp`` pipeline as ONE native dag-ml CV+refit run (S7).
REP FUSION is a one-time host RESHAPE feeding an already-native multimodal path:
* ``rep_to_sources`` — each replicate becomes a feature SOURCE, so the reshaped dataset is
MULTI-SOURCE; the envelope auto-emits a ``feature_block_set`` (S3 early fusion), and an MB-PLS
model takes the per-source block list (S5 intermediate fusion) — both already native.
* ``rep_to_pp`` — each replicate becomes a PROCESSING layer, so the reshaped dataset is
single-source with N stacked layers; the FLAT_2D materialization hstacks them by processing
order (the feature-axis concat S6 already does), which the legacy ``tabular_numeric`` path runs.
The reshaped dataset lives only in host memory (the on-disk dataset has no such structure), so it
is PICKLED for the adapter (the same mechanism ``sample_augmentation`` uses) — the adapter resolves
the exact reshaped sources/processings, identity-keyed by the physical sample_id.
LEAKAGE: after the reshape the unit of analysis is the physical SAMPLE (N reps → N sources/layers
of ONE sample), so folds/OOF are over SAMPLES — distinct from a plain repetition dataset (#21,
rep-grain). A sample's N source-blocks (or N processing layers) all ride ONE row and therefore the
SAME fold side by construction; per-source / per-layer preprocessing fits on fold-train only (the
per-block X-chain, like the multi-source path). No cross-sample mixing.
GENERATORS — an ``_or_`` / ``_cartesian_`` sweep INSIDE the rep-fusion body (e.g.
``[{"rep_to_pp": ...}, {"_or_": [SNV, MSC]}, splitter, model]``) is expanded in Python (the body
minus the reshape step via ``expand_spec``), each concrete variant runs through the reshaped
early-fusion path returning its raw native ScoreSet, and the per-variant projection (#55,
:func:`~nirs4all.pipeline.dagml.result._project_operator_sweep`) combines them into the FULL legacy
per-variant table — selecting the CV winner, projecting EVERY variant's CV rows (legacy
num_predictions parity) under its OWN ``config_name`` / ``model_name``, and refitting the winner
only. ``variant_config_names`` is the ordered legacy per-variant config names (derived by the caller
from the FULL pipeline — the reshape step IS part of the legacy config-name hash — in ``expand_spec``
order). A single (non-sweep) variant maps straight through the winner-only projection, byte-identical
to the pre-#55 shape.
"""
import pickle
from nirs4all.pipeline.config.generator import expand_spec
body = [step for step in pipeline if not _is_rep_fusion_step(step)]
variants = expand_spec(body)
run_dir.mkdir(parents=True, exist_ok=True)
variant_scores = [
_run_rep_fusion_concrete_scores(variant, rep_step, spectro, dataset_arg, cli, venv_python, run_dir / f"variant{index}", metric, pickle, random_state=random_state)
for index, variant in enumerate(variants)
]
if len(variant_scores) == 1:
scores, model_name, refit_artifacts = variant_scores[0]
return _scores_to_run_result(scores, spectro.name, model_name, metric, task_type, config_name=config_name, refit_artifacts=refit_artifacts)
# A sweep INSIDE the rep-fusion body: combine every reshaped-variant's ScoreSet into the full
# per-variant legacy table (#55) — same machinery the main operator-sweep path uses. `_project_operator_sweep`
# consumes score/model pairs, so split off the per-variant refit_artifacts and thread them as a separate by-index
# list (the projection persists the WINNER's model artifacts only).
sweep_scores = [(scores, model_name) for scores, model_name, _artifacts in variant_scores]
refit_artifacts_by_index = [artifacts for _scores, _model_name, artifacts in variant_scores]
return _project_operator_sweep(sweep_scores, spectro.name, metric, task_type, is_classification, variant_config_names or [], refit_artifacts_by_index=refit_artifacts_by_index)
def _run_rep_fusion_concrete_scores(body: Any, rep_step: dict[str, Any], spectro: Any, dataset_arg: str, cli: str, venv_python: str, run_dir: Path, metric: str, pickle: Any, random_state: int | None = None) -> tuple[dict[str, Any], str, list[dict[str, Any]]]:
"""One concrete rep-fusion variant: reshape a fresh dataset copy, run the sample-grain CV+refit, return raw scores.
Returns ``(scores, model_name, refit_artifacts)`` — the raw native ScoreSet + the model
label + the captured fitted REFIT estimators
(``outcome["refit_artifacts"]``, P3 Slice 2c-i) — so a sweep inside the body can COMBINE every variant's
ScoreSet into one per-variant projection (legacy num_predictions parity) AND persist the winner's model
artifacts, exactly like :func:`_run_concrete_scores` on the main path. The single-variant caller wraps
this through :func:`_scores_to_run_result` (winner-only projection).
"""
import copy
steps, splitter = _split_pipeline(body)
if splitter is None:
raise DagMlUnsupported("engine='dag-ml' requires a cross-validator step (e.g. KFold) in the pipeline")
_assert_supported_operators(steps)
steps = _apply_model_params(steps)
# Reshape a FRESH copy per variant so each variant's pickled dataset is independent (and the
# caller's spectro stays the original rep dataset, unmutated, for any later use).
reshaped = copy.deepcopy(spectro)
_reshape_for_rep_fusion(rep_step, reshaped)
identity = mint_identity(reshaped)
# The reshape leaves every physical sample in `partition: train` (the splitter follows the reshape),
# so the CV universe is the full reshaped sample set. Folds are sample-grain — a sample's N
# source-blocks / processing-layers ride ONE row, so they cannot split across the fold boundary.
pool = reshaped.index_column("sample", {"partition": "train"})
folds = _build_folds(splitter, reshaped, pool, set())
envelope = build_envelope(reshaped, identity, sample_ints=pool, group_by_sample=_split_group_grain(splitter, reshaped, pool))
dsl = assemble_cv_refit_dsl(steps, identity, envelope, folds, dsl_id="nirs4all-pipeline", n_splits=len(folds))
import dag_ml
dsl["root_seed"] = random_state if random_state is not None else 0
graph = dag_ml.compile_pipeline_dsl_artifact_with_controllers(dsl, controller_manifests()).graph.to_dict()
run_dir.mkdir(parents=True, exist_ok=True)
pickle_path = run_dir / "reshaped_dataset.pkl"
pickle_path.write_bytes(pickle.dumps(reshaped))
outcome = run_cv_refit_bundle(
dsl=dsl, envelope=envelope, graph=graph, dataset_path=dataset_arg, workdir=run_dir, dagml_cli=cli, venv_python=venv_python, selection_metric=metric, dataset_pickle=str(pickle_path), dataset=reshaped, random_state=random_state
)
if outcome["returncode"] != 0:
_raise_run_failure(outcome, "dag-ml rep-fusion run failed")
# Legacy labels the model by the model step's own name/class only — the rep_to_sources / rep_to_pp
# reshape is a dataset transform, NOT part of the model_name. Emit the bare `_model_name` (e.g.
# "PLSRegression"), matching legacy get_models() exactly (no "rep_to_sources_" / "rep_to_pp_" prefix).
return outcome["scores"], _model_name(steps), outcome["refit_artifacts"]
def _apply_sample_augmentation(aug_step: dict[str, Any], spectro: Any, context: Any | None = None) -> None:
"""Run nirs4all's REAL ``SampleAugmentationController`` to add synthetic train rows in place.
Reuses the production machinery (the controller delegates to ``TransformerMixinController``, which
fits the augmentation transformer on fold-train data and inserts each synthetic row via
``dataset.add_samples_batch`` with index ``{partition: train, origin: <base sample int>,
augmentation: <op>}``) — no augmentation logic is reimplemented here. A minimal real ``StepRunner``
+ ``RuntimeContext`` drive it (no workspace/artifacts), exactly as the orchestrator would. The
children land in ``partition: train`` with their ``origin`` set to the base sample, which is what
the base→child fold expansion and the envelope's augmentation grain key off.
"""
from nirs4all.controllers.data.sample_augmentation import SampleAugmentationController
from nirs4all.pipeline.config.context import DataSelector, ExecutionContext, PipelineState, RuntimeContext, StepMetadata
from nirs4all.pipeline.steps.parser import ParsedStep, StepType
from nirs4all.pipeline.steps.step_runner import StepRunner
step_info = ParsedStep(operator=None, keyword="sample_augmentation", step_type=StepType.DIRECT, original_step=aug_step, metadata={})
if context is None:
context = ExecutionContext(
selector=DataSelector(partition=None, processing=[["raw"]] * spectro.features_sources()),
state=PipelineState(),
metadata=StepMetadata(),
)
runtime_context = RuntimeContext()
runtime_context.step_runner = StepRunner(verbose=0, mode="train")
runtime_context.save_artifacts = False
runtime_context.save_charts = False
SampleAugmentationController().execute(step_info, spectro, context, runtime_context, mode="train")
def _augment_fold_train(aug_steps: list[dict[str, Any]], spectro: Any, fold_train: list[int], context: Any | None = None, chart_snapshots: list[Any] | None = None) -> list[tuple[int, np.ndarray | list[np.ndarray]]]:
"""Augment a fold's TRAIN only and return each child's source-preserving features.
A FRESH copy of ``spectro`` is restricted so ``partition: train`` is exactly ``fold_train`` (the
rest held out), then nirs4all's real augmentation machinery runs — so a STATEFUL/SUPERVISED/balanced
augmenter fits inside this fold's train ONLY (its neighbors / global mean / class balance never see
the fold's validation rows). The created children (origin in ``fold_train``) are read back as flat
feature rows, in creation order — one tuple per child, so an origin augmented multiple times (``count``
> 1 / a balancing factor) yields several. The caller inserts them into the master dataset and records
the fold→children map; the fold copy is discarded. Leakage-safe by construction: each fold's children
come only from its train.
"""
import copy
fold_ds = copy.deepcopy(spectro)
fold_ds._indexer.update_by_filter({"partition": "train"}, {"partition": "hold"}) # noqa: SLF001
fold_ds._indexer.update_by_indices(list(fold_train), {"partition": "train"}) # noqa: SLF001
fold_ds._invalidate_content_hash() # noqa: SLF001
before = {int(s) for s in fold_ds.index_column("sample", {})}
for aug_step in aug_steps:
_apply_sample_augmentation(aug_step, fold_ds, context)
if chart_snapshots is not None:
chart_snapshots.append(copy.deepcopy(fold_ds))
samples = [int(s) for s in fold_ds.index_column("sample", {})]
origins = [int(o) for o in fold_ds.index_column("origin", {})]
children: list[tuple[int, np.ndarray | list[np.ndarray]]] = []
for sample_int, origin_int in zip(samples, origins, strict=True):
if sample_int not in before and sample_int != origin_int:
children.append((origin_int, fold_ds.x_rows([sample_int], layout="3d", concat_source=False)))
return children
def _build_fold_local_children(aug_steps: list[dict[str, Any]], spectro: Any, base_folds: list[tuple[list[int], list[int]]], base_train: list[int], context: Any | None = None, chart_snapshots: list[Any] | None = None) -> tuple[dict[str, dict[int, list[int]]], dict[int, str]]:
"""Augment fold-by-fold + a full-train refit pass; insert all children into ``spectro`` in place.
For each fold (key ``"fold{i}"``, matching :func:`build_fold_set`'s fold ids) and the full-train
refit (key ``"refit"``) the augmenter is fit inside that train pool only (:func:`_augment_fold_train`),
and the resulting children are appended to the master ``spectro`` as ``partition: train`` rows with
their ``origin``. Returns ``(fold_children, augmentation_by_sample)`` where ``fold_children`` is
``{fold_label: {origin_int: [child_int, ...]}}`` (the resolver's fold-local expansion map) and
``augmentation_by_sample`` tags every inserted child for the envelope's augmentation metadata.
All folds' children coexist in one dataset and one base-grain envelope (each child shares its
origin's ``sample_id``, a base train id in the fold set → the origin-boundary contract holds), but
they are kept fold-distinct host-side via the returned map — a fold's children only ever join that
fold's fit-train (see :meth:`MaterializationResolver.expand_with_augmented_children`).
"""
transform_label = "+".join(_augmentation_label(step) for step in aug_steps)
passes: list[tuple[str, list[int]]] = [(f"fold{index}", train_ints) for index, (train_ints, _val) in enumerate(base_folds)]
passes.append(("refit", base_train))
fold_children: dict[str, dict[int, list[int]]] = {}
augmentation_by_sample: dict[int, str] = {}
base_spectro = copy.deepcopy(spectro)
for fold_label, fold_train in passes:
children = _augment_fold_train(aug_steps, base_spectro, fold_train, context, chart_snapshots if fold_label == "refit" else None)
if not children:
fold_children[fold_label] = {}
continue
first_child = children[0][1]
rows: np.ndarray | list[np.ndarray]
if isinstance(first_child, list):
rows = [np.concatenate([child_x[source] for _, child_x in children], axis=0) for source in range(len(first_child))]
else:
rows = np.concatenate([child_x for _, child_x in children], axis=0)
indexes = [{"partition": "train", "origin": origin_int, "augmentation": f"{transform_label}|{fold_label}"} for origin_int, _x in children]
before = {int(s) for s in spectro.index_column("sample", {})}
spectro.add_samples_batch(data=rows, indexes_list=indexes)
samples = [int(s) for s in spectro.index_column("sample", {})]
origins = [int(o) for o in spectro.index_column("origin", {})]
by_origin: dict[int, list[int]] = {}
for sample_int, origin_int in zip(samples, origins, strict=True):
if sample_int not in before and sample_int != origin_int:
by_origin.setdefault(origin_int, []).append(sample_int)
augmentation_by_sample[sample_int] = transform_label
fold_children[fold_label] = by_origin
return fold_children, augmentation_by_sample
def _build_fold_local_prefix_views(
prefix: list[Any], spectro: Any, base_folds: list[tuple[list[int], list[int]]],
base_train: list[int], context: Any | None = None, chart_snapshots: list[Any] | None = None,
chart_transform_snapshots: dict[tuple[int, int], Any] | None = None,
chart_transform_offset: int = 0,
) -> tuple[dict[str, dict[int, list[int]]], dict[int, str], dict[str, tuple[Any, dict[int, int], set[int]]], list[Any]]:
"""Run the ordered augmentation prefix in every train fold and in the refit pool.
The master dataset owns stable wire identities and targets. Each fold copy owns its fitted
preprocessing and synthetic spectra; ``local_ids`` translates master child ids back to the
corresponding child in that copy. No validation row participates in a fold's augmentation or
fitted preprocessing, though its transformed features remain available for prediction.
"""
label = "+".join(_augmentation_label(step) for step in prefix if _is_augmentation_step(step))
passes = [(f"fold{index}", train) for index, (train, _val) in enumerate(base_folds)]
passes.append(("refit", base_train))
fold_children: dict[str, dict[int, list[int]]] = {}
augmentation_by_sample: dict[int, str] = {}
feature_views: dict[str, tuple[Any, dict[int, int], set[int]]] = {}
replay_stages: list[Any] = []
base_spectro = copy.deepcopy(spectro)
for fold_label, fold_train in passes:
fold_ds = copy.deepcopy(base_spectro)
fold_ds._indexer.update_by_filter({"partition": "train"}, {"partition": "hold"}) # noqa: SLF001
fold_ds._indexer.update_by_indices(list(fold_train), {"partition": "train"}) # noqa: SLF001
fold_ds._invalidate_content_hash() # noqa: SLF001
before = {int(sample) for sample in fold_ds.index_column("sample", {})}
stages = _materialize_augmentation_prefix(
prefix, fold_ds, copy.deepcopy(context),
chart_snapshots if fold_label == "refit" else None,
chart_transform_snapshots if fold_label == "refit" else None,
transform_offset=chart_transform_offset,
)
if fold_label == "refit":
replay_stages = stages
samples = [int(sample) for sample in fold_ds.index_column("sample", {})]
origins = [int(origin) for origin in fold_ds.index_column("origin", {})]
local_origins = dict(zip(samples, origins, strict=True))
new_samples = [sample for sample in samples if sample not in before]
roots: list[int] = []
for child in new_samples:
root = local_origins[child]
visited = {child}
while root not in before:
if root in visited or root not in local_origins:
raise ValueError(f"invalid augmentation origin chain for sample {child}")
visited.add(root)
root = local_origins[root]
roots.append(root)
local_ids: dict[int, int] = {}
by_origin: dict[int, list[int]] = {}
if new_samples:
# Master rows are identity/target placeholders only; all feature requests for this
# regime are served from the selected fold copy by MaterializationResolver.
placeholder_rows = [spectro.x_rows([root], layout="3d", concat_source=False) for root in roots]
first = placeholder_rows[0]
rows = ([np.concatenate([row[source] for row in placeholder_rows], axis=0) for source in range(len(first))]
if isinstance(first, list) else np.concatenate(placeholder_rows, axis=0))
indexes = [{"partition": "train", "origin": root, "augmentation": f"{label}|{fold_label}"} for root in roots]
master_before = {int(sample) for sample in spectro.index_column("sample", {})}
spectro.add_samples_batch(data=rows, indexes_list=indexes)
master_new = [int(sample) for sample in spectro.index_column("sample", {}) if int(sample) not in master_before]
if len(master_new) != len(new_samples):
raise ValueError("fold-local augmentation changed child count while assigning identities")
for master_child, local_child, root in zip(master_new, new_samples, roots, strict=True):
local_ids[master_child] = local_child
by_origin.setdefault(root, []).append(master_child)
augmentation_by_sample[master_child] = label
fold_children[fold_label] = by_origin
excluded_local = {int(sample) for sample in fold_ds.index_column("sample", {"excluded": True})}
excluded_master = (excluded_local & before) | {master for master, local in local_ids.items() if local in excluded_local}
feature_views[fold_label] = (fold_ds, local_ids, excluded_master)
return fold_children, augmentation_by_sample, feature_views, replay_stages
def _apply_pre_augmentation_steps(pre_aug_steps: list[Any], spectro: Any, context: Any | None = None, chart_capture: Callable[[Any, Any], None] | None = None) -> tuple[Any, list[Any]]:
"""Materialize preceding transforms/exclusions and capture transform replay in order."""
from nirs4all.pipeline.config.context import DataSelector, ExecutionContext, PipelineState, RuntimeContext, StepMetadata
from nirs4all.pipeline.steps.step_runner import StepRunner
if context is None:
context = ExecutionContext(
selector=DataSelector(partition=None, processing=[["raw"]] * spectro.features_sources()),
state=PipelineState(),
metadata=StepMetadata(),
)
if not pre_aug_steps:
return context, []
_assert_supported_operators(pre_aug_steps)
replay_stages: list[Any] = []
runner = StepRunner(verbose=0, mode="train")
runtime_context = RuntimeContext()
runtime_context.step_runner = runner
runtime_context.save_artifacts = False
runtime_context.save_charts = False
index = 0
while index < len(pre_aug_steps):
step = pre_aug_steps[index]
from .detect import _simple_duplication_merge_mode
branch_pair = (
isinstance(step, dict) and "branch" in step
and index + 1 < len(pre_aug_steps)
and _simple_duplication_merge_mode(pre_aug_steps[index + 1]) == "features"
)
replay_steps = pre_aug_steps[index:index + 2] if branch_pair else [step]
changes_x = not (isinstance(step, dict) and any(key in step for key in ("exclude", "tag", "y_processing")))
raw_train = np.asarray(spectro.x({"partition": "train"}, layout="2d", include_augmented=True)).copy() if changes_x else None
replay_stage = _capture_pre_augmentation_replay(replay_steps, spectro) if changes_x else None
for replay_step in replay_steps:
result = runner.execute(replay_step, spectro, context, runtime_context, prediction_store=None)
context = result.updated_context
runtime_context.step_number += 1
_verify_pre_augmentation_replay(replay_stage, spectro, raw_train)
if replay_stage is not None:
replay_stages.append(replay_stage)
if changes_x and chart_capture is not None:
chart_capture(replay_steps[-1], spectro)
index += len(replay_steps)
return context, replay_stages
class _SourcePreAugmentationReplay(TransformerMixin, BaseEstimator):
"""Replay fitted pre-augmentation transforms on each raw source block."""
def __init__(self, chains: list[Any], input_widths: list[int]) -> None:
self.chains = chains
self.input_widths = input_widths
def fit(self, X: np.ndarray, y: Any = None) -> _SourcePreAugmentationReplay:
"""Refit the per-source chains when a captured predictor is retrained."""
start = 0
for width, chain in zip(self.input_widths, self.chains, strict=True):
chain.fit(np.asarray(X)[:, start:start + width], y)
start += width
return self
def transform(self, X: np.ndarray) -> np.ndarray:
"""Split the raw fusion matrix, transform sources, then concatenate."""
x = np.asarray(X)
if x.shape[1] != sum(self.input_widths):
raise ValueError("pre-augmentation source widths differ from the fitted layout")
blocks = []
start = 0
for width, chain in zip(self.input_widths, self.chains, strict=True):
block = x[:, start:start + width]
for _, transformer in chain.steps:
block = np.asarray(transformer.transform(block))
blocks.append(block)
start += width
return np.concatenate(blocks, axis=1)
def _capture_pre_augmentation_replay(pre_aug_steps: list[Any], spectro: Any) -> Any | None:
"""Fit the same prefix on raw training rows for captured predictor replay."""
if not pre_aug_steps:
return None
from sklearn.base import clone
from sklearn.pipeline import make_pipeline
from nirs4all.operators.transforms.concat import FeatureConcat
from nirs4all.pipeline.dagml.node_runner import _CoordinateTransform
from nirs4all.pipeline.dagml.steps import _needs_wavelength_injection
from nirs4all.pipeline.dagml_bridge import _lower_feature_augmentation
_assert_supported_operators(pre_aug_steps)
branch_merge = len(pre_aug_steps) == 2 and isinstance(pre_aug_steps[0], dict) and "branch" in pre_aug_steps[0]
if branch_merge:
from .detect import _duplication_branch_bodies, _selected_duplication_feature_branches, _simple_duplication_merge_mode
branches = _duplication_branch_bodies(pre_aug_steps[0])
mode = _simple_duplication_merge_mode(pre_aug_steps[1])
if branches is None or mode != "features":
raise DagMlUnsupported("pre-augmentation branch replay requires a duplication feature merge")
branches = _selected_duplication_feature_branches(branches, pre_aug_steps[1])
if branches is None:
raise DagMlUnsupported("pre-augmentation branch replay has an invalid feature selection")
try:
# A wrapped preprocessing step carries fit-scope metadata, not estimator
# parameters. Mirror the legacy controller's fit cohort while keeping the
# full feature chain available for prediction replay.
fit_on_all = any(isinstance(step, dict) and step.get("fit_on_all") is True
for step in pre_aug_steps)
fit_selector = {} if fit_on_all else {"partition": "train"}
y = np.asarray(spectro.y(fit_selector, include_augmented=True))
raw_blocks = spectro.x(fit_selector, layout="2d", concat_source=False, include_augmented=True)
blocks = raw_blocks if isinstance(raw_blocks, list) else [raw_blocks]
chains = []
for source_index, block in enumerate(blocks):
transforms = (
[_branch_merge_transformer_step(cast(list[list[Any]], branches), "features")]
if branch_merge else [
FeatureConcat(**_lower_feature_augmentation(step)["params"])
if isinstance(step, dict) and "feature_augmentation" in step else
clone(step["preprocessing"] if isinstance(step, dict) and "preprocessing" in step else step)
for step in pre_aug_steps
]
)
transforms = [
_CoordinateTransform(transform, tuple(str(value) for value in spectro.wavelengths_cm1(source_index)), source_index)
if _needs_wavelength_injection(transform) else transform
for transform in transforms
]
chain = make_pipeline(*transforms)
chain.fit(np.asarray(block), y[:, 0] if y.ndim > 1 else y)
chains.append(chain)
except (TypeError, AttributeError) as exc:
raise DagMlUnsupported(f"cannot capture pre-augmentation transform replay: {exc}") from exc
if len(chains) == 1:
return chains[0]
return _SourcePreAugmentationReplay(chains, [np.asarray(block).shape[1] for block in blocks])
def _attach_pre_augmentation_replay(result: RunResult, stages: list[Any]) -> RunResult:
"""Prepend the fitted, verified prefix to each captured final model."""
if not stages:
return result
from sklearn.pipeline import Pipeline
pre_steps = [
(f"pre_augmentation_{index}", stage if isinstance(stage, _SourcePreAugmentationReplay) else stage.steps[0][1])
for index, stage in enumerate(stages)
]
for artifact in result._dagml_refit_artifacts: # noqa: SLF001 - captured host artifact contract
artifact["estimator"] = Pipeline([*pre_steps, ("model", artifact["estimator"])])
return result
def _verify_pre_augmentation_replay(chain: Any | None, spectro: Any, raw_train: np.ndarray | None) -> None:
"""Refuse an export prefix that differs from the data the controller materialized."""
if chain is None or raw_train is None:
return
actual = np.asarray(spectro.x({"partition": "train"}, layout="2d", include_augmented=True))
if isinstance(chain, _SourcePreAugmentationReplay):
expected = chain.transform(raw_train)
else:
expected = raw_train
for _, transformer in chain.steps:
expected = np.asarray(transformer.transform(expected))
# Fitting a float32 projection twice can differ by a few ulps near zero even when the
# production controller and export replay use the same training rows and operator.
if actual.shape != expected.shape or not np.allclose(actual, expected, rtol=1e-5, atol=1e-5):
raise DagMlUnsupported("pre-augmentation transform replay differs from the fitted controller output")
def _augmentation_grain(spectro: Any, transform_label: str) -> tuple[list[int], dict[int, str]]:
"""Post-augmentation grain: ``(augmented_ints, augmentation_by_sample)``.
``augmented_ints`` is every augmented child (``sample != origin``); ``augmentation_by_sample`` tags
each with the augmentation transform id for the envelope's structured ``augmentation`` metadata. The
fit-time base→child expansion is recomputed by the resolver from the minted identity, not here.
"""
samples = [int(s) for s in spectro.index_column("sample", {})]
origins = [int(o) for o in spectro.index_column("origin", {})]
augmented_ints = [sample_int for sample_int, origin_int in zip(samples, origins, strict=True) if sample_int != origin_int]
augmentation_by_sample = dict.fromkeys(augmented_ints, transform_label)
return augmented_ints, augmentation_by_sample
def _augmentation_label(aug_step: dict[str, Any]) -> str:
"""A stable transform label for the augmentation step's structured envelope metadata."""
transformers = _augmentation_transformers(aug_step)
return "+".join(type(transformer).__name__ for transformer in transformers) or "sample_augmentation"
def _augmentation_transformers(aug_step: dict[str, Any]) -> list[Any]:
"""The deserialized transformer instances of an augmentation step (mirrors the controller).
A transformer may be given bare or wrapped in a ``{"transformer": op, ...}`` dict (per-transformer
variation_scope); both forms resolve to the instance via ``deserialize_component``, exactly as
:meth:`SampleAugmentationController.execute` parses them.
"""
from nirs4all.controllers.transforms.transformer import TransformerMixinController
from nirs4all.pipeline.config.component_serialization import deserialize_component
raw = aug_step["sample_augmentation"].get("transformers", [])
transformers = [deserialize_component(t["transformer"] if isinstance(t, dict) and "transformer" in t else t) for t in raw]
# Refuse X-only mixup before statelessness probing or row materialization.
for transformer in transformers:
TransformerMixinController.validate_target_preserving_operator(transformer)
return transformers
def _operator_is_stateless(operator: Any) -> bool:
"""Whether an augmentation transformer learns NO data-dependent state in ``fit``.
The first augmentation slice augments ONCE globally (before folds exist), so a transformer that
fits on the whole train partition would see future fold-validation rows. That is leakage-free ONLY
for STATELESS per-sample augmenters (Gaussian/multiplicative noise, scatter, baseline, spline,
wavelength warps — ``fit`` only seeds an RNG); a STATEFUL/SUPERVISED one (mixup storing neighbors,
a global-mean reference, a supervised transform) leaks. There is no declared marker on these
operators, so the signal is twofold and conservative:
* **supervised** — a ``requires_y`` tag (``_more_tags()``/``_tags``) means ``fit`` consumes y; reject.
* **fit-learned data state** — fit the transformer on two differently-distributed dummy datasets and
compare its post-fit attributes (sklearn fitted attrs ``*_`` + any ndarray, excluding the seeded
RNG). A transformer whose state varies with the fit data carries learned state (``X_fit_``,
``global_mean_``, …) → stateful → reject. A clone is fit each time so no instance state is shared;
any error during the probe is treated as NOT stateless (fail closed).
"""
from nirs4all.controllers.transforms.transformer import TransformerMixinController
if TransformerMixinController._requires_y(operator): # noqa: SLF001 - reuse the supervised-tag check
return False
try:
from sklearn.base import clone
rng = np.random.default_rng(0)
probe_a = rng.normal(size=(24, 32))
probe_b = rng.normal(loc=8.0, scale=4.0, size=(24, 32))
def _fit_state(data: np.ndarray) -> dict[str, Any]:
fitted = clone(operator).fit(data)
return {name: value for name, value in vars(fitted).items() if (name.endswith("_") or isinstance(value, np.ndarray)) and not name.startswith("_")}
state_a, state_b = _fit_state(probe_a), _fit_state(probe_b)
for name in set(state_a) | set(state_b):
left, right = state_a.get(name), state_b.get(name)
if isinstance(left, np.ndarray) or isinstance(right, np.ndarray):
if not (left is not None and right is not None and np.array_equal(np.asarray(left), np.asarray(right))):
return False
elif left != right:
return False
except Exception: # noqa: BLE001 - any probe failure ⇒ cannot prove stateless ⇒ fail closed
return False
return True
def _augmentation_is_leakage_free(aug_step: dict[str, Any]) -> bool:
"""Whether a ``sample_augmentation`` step is safe to run as ONE GLOBAL pre-fold augmentation.
True (→ the GLOBAL path, #8) ONLY when neither leakage vector applies to global augmentation:
* the **balanced/supervised mode** (a ``balance`` key) is NOT used — it fits class/bin targets on
the whole train y, so the synthetic counts depend on the (future-fold-val-inclusive) label set;
* EVERY transformer is stateless (:func:`_operator_is_stateless`).
False (→ the FOLD-LOCAL path, #32) for a stateful/supervised/balanced augmentation: it is fit inside
each fold's train only (and a full-train refit pass), so it never sees a fold's validation rows. Both
paths are leakage-safe; this only routes which one :func:`_run_augmentation` uses.
"""
config = aug_step["sample_augmentation"]
if "balance" in config:
return False
transformers = _augmentation_transformers(aug_step)
return bool(transformers) and all(_operator_is_stateless(transformer) for transformer in transformers)
def _materialize_augmentation_prefix(
prefix: list[Any], spectro: Any, context: Any | None = None,
chart_snapshots: list[Any] | None = None,
chart_transform_snapshots: dict[tuple[int, int], Any] | None = None,
*, transform_offset: int = 0,
) -> list[Any]:
"""Execute transforms and augmentation in public order, capturing prediction replay."""
replay_stages: list[Any] = []
pending: list[Any] = []
augmentation_index = 0
transform_index = transform_offset
def capture_transform(_step: Any, dataset: Any) -> None:
nonlocal transform_index
transform_index += 1
if chart_transform_snapshots is not None:
chart_transform_snapshots[(augmentation_index, transform_index)] = copy.deepcopy(dataset)
for step in prefix:
if _is_augmentation_step(step):
context, stages = _apply_pre_augmentation_steps(pending, spectro, context, capture_transform)
replay_stages.extend(stages)
pending = []
_apply_sample_augmentation(step, spectro, context)
augmentation_index += 1
if chart_snapshots is not None:
chart_snapshots.append(copy.deepcopy(spectro))
else:
pending.append(step)
if pending:
_context, stages = _apply_pre_augmentation_steps(pending, spectro, context, capture_transform)
replay_stages.extend(stages)
return replay_stages
def _post_augmentation_exclusion_prefix_length(steps: list[Any]) -> int:
"""How much of a post-augmentation transform/exclude prefix runs before modeling."""
from .detect import _is_exclude_step
transform_end = 0
for step in steps:
is_transform = (
isinstance(step, dict) and "preprocessing" in step
) or (not isinstance(step, dict) and hasattr(step, "fit") and hasattr(step, "transform")
and not hasattr(step, "predict"))
if not (is_transform or _is_exclude_step(step)):
break
transform_end += 1
exclude_indices = [index for index, step in enumerate(steps) if _is_exclude_step(step)]
if not exclude_indices:
return transform_end
end = exclude_indices[-1] + 1
if any(_is_split_step(step) or (isinstance(step, dict) and any(key in step for key in ("model", "branch", "merge"))) for step in steps[:end]):
raise DagMlUnsupported("exclude after a splitter, model or branch cannot be materialized before augmentation training")
return max(end, transform_end)
def _run_augmentation_full_train(
pipeline: list[Any], spectro: Any, dataset_arg: str, cli: str,
venv_python: str, run_dir: Path, *, metric: str, task_type: str,
config_name: str, random_state: int | None = None,
train_sample_ids: list[int] | None = None,
) -> RunResult:
"""Apply sample augmentation before a DAG-owned full-training phase."""
from .full_train import run_full_train
aug_indices = [index for index, step in enumerate(pipeline) if _is_augmentation_step(step)]
from .detect import _is_exclude_step
pre_aug_steps = pipeline[:aug_indices[0]]
early_models = [step for step in pre_aug_steps if isinstance(step, dict) and "model" in step]
next_model = next((index for index in range(aug_indices[-1] + 1, len(pipeline))
if isinstance(pipeline[index], dict) and "model" in pipeline[index]), len(pipeline))
if early_models and (
pre_aug_steps[-len(early_models):] != early_models
or len(aug_indices) > 1
or aug_indices != list(range(aug_indices[0], aug_indices[-1] + 1))
or any(not _is_split_step(step) for step in pipeline[aug_indices[-1] + 1:next_model])
or any(_is_exclude_step(step) for step in pipeline[aug_indices[-1] + 1:])
):
return _run_interleaved_full_train_checkpoints(
pipeline, spectro, dataset_arg, cli, venv_python, run_dir,
metric=metric, task_type=task_type, config_name=config_name,
random_state=random_state, train_sample_ids=train_sample_ids,
)
chart_snapshots: list[Any] | None = [] if getattr(spectro, "_dagml_capture_aug_charts", False) else None
chart_transform_snapshots: dict[tuple[int, int], Any] | None = {} if chart_snapshots is not None else None
after_aug = aug_indices[-1] + 1
materialize_end = after_aug + _post_augmentation_exclusion_prefix_length(pipeline[after_aug:])
pre_aug_steps = pipeline[:aug_indices[0]]
checkpoint_steps = [step for step in pre_aug_steps if isinstance(step, dict) and "model" in step]
if checkpoint_steps:
if (len(aug_indices) != 1 or materialize_end != after_aug
or pre_aug_steps[-len(checkpoint_steps):] != checkpoint_steps
or not any(isinstance(step, dict) and "model" in step for step in pipeline[after_aug:])):
raise DagMlUnsupported("sequential full-training checkpoints across this augmentation shape need distinct native fit views")
materialize_prefix = [*pre_aug_steps[:-len(checkpoint_steps)], *pipeline[aug_indices[0]:materialize_end]]
else:
materialize_prefix = pipeline[:materialize_end]
replay_stages = _materialize_augmentation_prefix(
materialize_prefix, spectro,
chart_snapshots=chart_snapshots, chart_transform_snapshots=chart_transform_snapshots,
)
# The host model still needs the Y transform: materializing the prefix for the
# augmentation controller does not add a Y node to the native model graph.
y_prefix_steps = [step for step in pipeline[:aug_indices[0]] if isinstance(step, dict) and "y_processing" in step]
post_aug_steps = [*y_prefix_steps, *checkpoint_steps, *pipeline[materialize_end:]]
from .detect import _detect_duplication_branch
duplication = _detect_duplication_branch(post_aug_steps)
if duplication is not None and duplication[1] == "mean":
post_aug_steps = [_branch_fusion_model_step(duplication[0], duplication[1], task_type)]
import pickle
run_dir.mkdir(parents=True, exist_ok=True)
pickle_path = run_dir / "augmented_dataset.pkl"
pickle_path.write_bytes(pickle.dumps(spectro))
result = run_full_train(
post_aug_steps, spectro, metric=metric, task_type=task_type,
config_name=config_name, augmented_train=True,
cli=cli, venv_python=venv_python, dataset_path=dataset_arg,
dataset_pickle=str(pickle_path), workdir=run_dir / "refit",
random_state=random_state, train_sample_ids=train_sample_ids,
base_fit_model_count=len(checkpoint_steps),
)
result = _attach_pre_augmentation_replay(result, replay_stages)
result._dagml_chart_aug_snapshots = chart_snapshots # type: ignore[attr-defined]
result._dagml_chart_transform_snapshots = chart_transform_snapshots # type: ignore[attr-defined]
return result
def _run_interleaved_full_train_checkpoints(
pipeline: list[Any], spectro: Any, dataset_arg: str, cli: str,
venv_python: str, run_dir: Path, *, metric: str, task_type: str,
config_name: str, random_state: int | None,
train_sample_ids: list[int] | None,
) -> RunResult:
"""Project independent unsplit checkpoint campaigns into one public result."""
import inspect
import pickle
import dag_ml
from nirs4all.data.predictions import Predictions
from .full_train import run_full_train
model_indices = [index for index, step in enumerate(pipeline) if isinstance(step, dict) and "model" in step]
if len(model_indices) < 2:
raise DagMlUnsupported("interleaved full-training checkpoints require at least two models")
campaigns: list[RunResult] = []
for candidate_index, model_index in enumerate(model_indices):
candidate = [
step for index, step in enumerate(pipeline[:model_index + 1])
if index == model_index or not (isinstance(step, dict) and "model" in step)
]
candidate_spectro = copy.deepcopy(spectro)
candidate_dir = run_dir / f"checkpoint{candidate_index}"
candidate_dir.mkdir(parents=True, exist_ok=True)
if any(_is_augmentation_step(step) for step in candidate):
campaigns.append(_run_augmentation_full_train(
candidate, candidate_spectro, dataset_arg, cli, venv_python,
candidate_dir, metric=metric, task_type=task_type,
config_name=config_name, random_state=random_state,
train_sample_ids=train_sample_ids,
))
else:
pickle_path = candidate_dir / "checkpoint_dataset.pkl"
pickle_path.write_bytes(pickle.dumps(candidate_spectro))
campaigns.append(run_full_train(
candidate, candidate_spectro, metric=metric, task_type=task_type,
config_name=config_name, cli=cli, venv_python=venv_python,
dataset_path=dataset_arg, dataset_pickle=str(pickle_path),
workdir=candidate_dir / "refit", random_state=random_state,
train_sample_ids=train_sample_ids,
))
candidate_scores = []
for index, campaign in enumerate(campaigns):
reports = cast(dict[str, Any], campaign._dagml_score_set)["reports"]
test_report = next((report for report in reports if report["partition"] == "test" and report["level"] == "sample"), None)
fit_report = next((report for report in reports if report["partition"] == "final" and report["level"] == "sample"), None)
report = test_report or fit_report
if report is None:
raise ValueError("full-training checkpoint lacks a native score report")
candidate_scores.append({"candidate_id": str(index), "metrics": {metric: float(report["metrics"][metric])}})
decision = dag_ml.select_candidate(
{"id": "select:interleaved_full_train_checkpoints", "metric": {
"name": metric, "objective": "maximize" if is_higher_better(metric) else "minimize",
}},
candidate_scores,
)
selected_index = int(decision["selected_candidate_id"])
accepted_fields = set(inspect.signature(Predictions.add_prediction).parameters) - {"self"}
predictions = Predictions()
for campaign in campaigns:
for row in campaign.predictions.filter_predictions():
predictions.add_prediction(**{key: value for key, value in row.items() if key in accepted_fields})
predictions.flush()
selected = campaigns[selected_index]
result = RunResult(predictions=predictions, per_dataset=copy.deepcopy(selected.per_dataset))
result._dagml_score_set = selected._dagml_score_set
result._dagml_checkpoint_score_sets = [campaign._dagml_score_set for campaign in campaigns] # type: ignore[attr-defined]
result._dagml_node_results = selected._dagml_node_results
result._dagml_refit_artifacts = selected._dagml_refit_artifacts
chart_source = next((campaign for campaign in reversed(campaigns) if hasattr(campaign, "_dagml_chart_aug_snapshots")), None)
if chart_source is not None:
result._dagml_chart_aug_snapshots = chart_source._dagml_chart_aug_snapshots # type: ignore[attr-defined]
result._dagml_chart_transform_snapshots = chart_source._dagml_chart_transform_snapshots # type: ignore[attr-defined]
return result
def _run_interleaved_augmentation_checkpoints(
pipeline: list[Any], spectro: Any, dataset_arg: str, cli: str,
venv_python: str, run_dir: Path, metric: str, task_type: str,
config_name: str, random_state: int | None, refit: bool,
) -> RunResult:
"""Run independent legacy checkpoints as separate native campaigns.
A transform or exclusion between checkpoints changes the data seen by the
later model. In particular an exclusion can change its validation FoldSet.
The native campaign has one FoldSet, so each prefix gets its own native
graph and envelope; core SELECT chooses the exported checkpoint.
"""
import pickle
import dag_ml
from .detect import _is_exclude_step
from .result import _variant_cv_score
model_indices = [index for index, step in enumerate(pipeline) if isinstance(step, dict) and "model" in step]
if len(model_indices) < 2:
raise DagMlUnsupported("interleaved augmentation checkpoints require at least two models")
campaigns: list[tuple[dict[str, Any], str, list[dict[str, Any]], Any, list[dict[str, Any]], dict[str, Any]]] = []
chart_source: RunResult | None = None
for candidate_index, model_index in enumerate(model_indices):
# Legacy's ordinary model steps are independent checkpoints. Preserve
# every data-changing predecessor, omitting only earlier model steps.
candidate = [
step for index, step in enumerate(pipeline[:model_index + 1])
if index == model_index or not (isinstance(step, dict) and "model" in step)
]
splitter_steps = [step for step in candidate if _is_split_step(step)]
if len(splitter_steps) != 1:
raise DagMlUnsupported("interleaved checkpoints require one shared public splitter")
original_folds = None
# An exclusion after augmentation must finish before the later model's
# splitter. Keep the public splitter's ORIGINAL fold assignment and
# filter excluded ids within each side, as legacy does when its splitter
# ran before the exclusion step.
if any(_is_exclude_step(step) for step in candidate):
_, splitter = _split_pipeline(candidate)
original_folds = _build_folds(
splitter, spectro, _split_base_samples(spectro), set(),
)
candidate = [step for step in candidate if not _is_split_step(step)]
candidate.insert(-1, splitter_steps[0])
candidate_spectro = copy.deepcopy(spectro)
candidate_dir = run_dir / f"checkpoint{candidate_index}"
candidate_dir.mkdir(parents=True, exist_ok=True)
if any(_is_augmentation_step(step) for step in candidate):
capture: dict[str, Any] = {}
candidate_result = _run_augmentation(
candidate, candidate_spectro, dataset_arg, cli, venv_python,
candidate_dir, metric, task_type, config_name=config_name,
random_state=random_state, capture=capture, refit=refit,
base_folds_override=original_folds,
)
campaigns.append((
capture["scores"], capture["model_name"], capture["results"],
capture["identity"], capture["refit_artifacts"], lane_record(HOST_CALLBACK_LANE),
))
chart_source = candidate_result
else:
pickle_path = candidate_dir / "checkpoint_dataset.pkl"
pickle_path.write_bytes(pickle.dumps(candidate_spectro))
campaigns.append(_run_concrete_scores(
candidate, candidate_spectro, dataset_arg, cli, venv_python,
candidate_dir, dataset_pickle=str(pickle_path),
random_state=random_state, refit=refit, metric=metric,
))
decision = dag_ml.select_candidate(
{"id": "select:interleaved_augmentation_checkpoints", "metric": {
"name": metric, "objective": "maximize" if is_higher_better(metric) else "minimize",
}},
[{"candidate_id": str(index), "metrics": {metric: _variant_cv_score(scores, metric)}}
for index, (scores, *_rest) in enumerate(campaigns)],
)
result = _project_operator_sweep(
[(scores, model_name) for scores, model_name, *_rest in campaigns],
spectro.name, metric, task_type, task_type != "regression",
[config_name] * len(campaigns),
results_by_index=[results for _scores, _name, results, _identity, _artifacts, _lane in campaigns],
identities_by_index=[identity for _scores, _name, _results, identity, _artifacts, _lane in campaigns],
refit_artifacts_by_index=[artifacts for _scores, _name, _results, _identity, artifacts, _lane in campaigns],
selected_index=int(decision["selected_candidate_id"]), emit_all_refits=True,
)
record_execution_lane(result, merge_lane_records([lane for *_rest, lane in campaigns]))
if chart_source is not None:
result._dagml_chart_aug_snapshots = chart_source._dagml_chart_aug_snapshots # type: ignore[attr-defined]
result._dagml_chart_transform_snapshots = chart_source._dagml_chart_transform_snapshots # type: ignore[attr-defined]
return result
def _run_augmentation(pipeline: list[Any], spectro: Any, dataset_arg: str, cli: str, venv_python: str, run_dir: Path, metric: str, task_type: str, config_name: str = "", random_state: int | None = None, capture: dict[str, Any] | None = None, refit: bool = True, base_folds_override: list[tuple[list[int], list[int]]] | None = None) -> RunResult:
"""Run a ``sample_augmentation`` pipeline as ONE native dag-ml CV+refit on augmented train.
Adds the synthetic train rows (real augmentation machinery), builds BASE-grain folds (each base
val sample validated once — a clean OOF partition that dag-ml/dag-ml-data accept) and a CV-universe
envelope (base + augmented children; observation-grain relations carrying each child's origin +
augmentation, deduped to the origin's sample grain in the schema). The fold-train views stay
base-grain + ``include_augmented_train`` so the host expands each base id to base + its children at
fit time — the children TRAIN, the OOF/validation/test never see them. The augmented dataset is
pickled for the adapter (augmentation is stochastic, not reproducible cross-process).
Two augmentation regimes, picked by :func:`_augmentation_is_leakage_free`:
* **GLOBAL** (stateless per-sample augmenter) — augment ONCE on the whole train; the children are
shared across folds. Leakage-free because the augmenter learns no data state (#8).
* **FOLD-LOCAL** (stateful / supervised / balanced augmenter) — augment inside EACH fold's train
(plus a full-train refit pass) separately, so each fold has its own children fit only on that
fold's train. A ``fold_children`` map keys the resolver's per-fold expansion; it is pickled with
the dataset so the adapter expands the right children per fold (#32). This is what makes the
stateful case (mixup neighbors, global-mean scatter, class balancing) leakage-safe.
Ordered transforms/exclusions around augmentation are materialized before the native phase;
duplicate feature/fusion branches and by-metadata separation use their native model paths.
Interleaved fold-local augmentation uses separate, ordered feature views for every fold and refit.
"""
import pickle
chart_snapshots: list[Any] | None = [] if getattr(spectro, "_dagml_capture_aug_charts", False) else None
chart_transform_snapshots: dict[tuple[int, int], Any] | None = {} if chart_snapshots is not None else None
from .detect import _is_exclude_step
aug_indices = [index for index, step in enumerate(pipeline) if _is_augmentation_step(step)]
pre_aug_steps = pipeline[:aug_indices[0]]
early_models = [step for step in pre_aug_steps if isinstance(step, dict) and "model" in step]
next_model = next((index for index in range(aug_indices[-1] + 1, len(pipeline))
if isinstance(pipeline[index], dict) and "model" in pipeline[index]), len(pipeline))
if early_models and (
pre_aug_steps[-len(early_models):] != early_models
or len(aug_indices) > 1
or aug_indices != list(range(aug_indices[0], aug_indices[-1] + 1))
or any(not _is_split_step(step) for step in pipeline[aug_indices[-1] + 1:next_model])
or any(_is_exclude_step(step) for step in pipeline[aug_indices[-1] + 1:])
):
return _run_interleaved_augmentation_checkpoints(
pipeline, spectro, dataset_arg, cli, venv_python, run_dir,
metric, task_type, config_name, random_state, refit,
)
# Legacy accepts the splitter before augmentation. Native folds are built on
# base sample ids independently of its position, so route it after the last
# augmentation while preserving the order of all data-changing steps.
last_aug = max(index for index, step in enumerate(pipeline) if _is_augmentation_step(step))
early_splitters = [index for index, step in enumerate(pipeline[:last_aug]) if _is_split_step(step)]
if early_splitters:
if len(early_splitters) != 1 or any(_is_split_step(step) for step in pipeline[last_aug + 1:]):
raise DagMlUnsupported("sample augmentation with multiple splitters needs an explicit fold policy")
pipeline = list(pipeline)
splitter_step = pipeline.pop(early_splitters[0])
last_aug = max(index for index, step in enumerate(pipeline) if _is_augmentation_step(step))
ordered_post_aug = _post_augmentation_exclusion_prefix_length(pipeline[last_aug + 1:])
pipeline.insert(last_aug + 1 + ordered_post_aug, splitter_step)
aug_indices = [index for index, step in enumerate(pipeline) if _is_augmentation_step(step)]
interleaved = aug_indices != list(range(aug_indices[0], aug_indices[-1] + 1))
aug_steps = [pipeline[index] for index in aug_indices]
fold_local = not all(_augmentation_is_leakage_free(step) for step in aug_steps)
aug_index, after_aug = aug_indices[0], aug_indices[-1] + 1
from .detect import _is_exclude_step
pre_aug_steps = pipeline[:aug_index]
checkpoint_steps = [step for step in pre_aug_steps if isinstance(step, dict) and "model" in step]
if checkpoint_steps:
if (not any(isinstance(step, dict) and "model" in step for step in pipeline[after_aug:])
or fold_local or interleaved
or pre_aug_steps[-len(checkpoint_steps):] != checkpoint_steps):
raise DagMlUnsupported("sequential model checkpoints across this augmentation shape need distinct native fit views")
pre_aug_steps = pre_aug_steps[:-len(checkpoint_steps)]
chart_transform_offset = sum(
(isinstance(step, dict) and set(step) == {"preprocessing"})
or (not isinstance(step, dict) and hasattr(step, "transform") and not hasattr(step, "predict"))
for step in pre_aug_steps
)
y_prefix_steps = [step for step in pre_aug_steps if isinstance(step, dict) and "y_processing" in step]
late_exclusion_length = _post_augmentation_exclusion_prefix_length(pipeline[after_aug:])
post_aug_steps = pipeline[after_aug + late_exclusion_length:]
materialized_early = bool((interleaved or late_exclusion_length) and not fold_local)
if checkpoint_steps and materialized_early:
raise DagMlUnsupported("sequential model checkpoints with ordered exclusion need distinct native fit views")
if materialized_early:
replay_stages = _materialize_augmentation_prefix(
pipeline[:after_aug + late_exclusion_length], spectro,
chart_snapshots=chart_snapshots, chart_transform_snapshots=chart_transform_snapshots,
)
prefix_context = None
else:
prefix_context, replay_stages = _apply_pre_augmentation_steps(pre_aug_steps, spectro)
steps, splitter = _split_pipeline(post_aug_steps)
if splitter is None:
raise DagMlUnsupported("engine='dag-ml' requires a cross-validator step (e.g. KFold) in the pipeline")
steps = [*y_prefix_steps, *checkpoint_steps, *steps]
if any(_is_exclude_step(step) for step in steps):
from .exclude import _resolve_exclude
first_non_exclude = next((index for index, step in enumerate(steps) if not _is_exclude_step(step)), len(steps))
if any(_is_exclude_step(step) for step in steps[first_non_exclude:]):
raise DagMlUnsupported("exclude after post-augmentation transforms needs ordered transform materialization")
steps, allowed, marked = _resolve_exclude(steps, spectro)
allowed_base = set(allowed) - marked
else:
allowed_base = None
from .detect import _detect_duplication_branch
duplication = _detect_duplication_branch(steps)
if duplication is not None:
branches, merge_mode = duplication
if merge_mode == "features":
model_step = next(step for step in steps if isinstance(step, dict) and "model" in step)
steps = [_branch_merge_transformer_step(branches, merge_mode), model_step]
elif merge_mode == "mean":
steps = [_branch_fusion_model_step(branches, merge_mode, task_type)]
from .detect import _detect_separation_branch
separation = _detect_separation_branch(steps)
if separation is not None:
if allowed_base is not None:
raise DagMlUnsupported("separation branch with post-augmentation exclusion needs branch-scoped exclusion views")
else:
_assert_supported_operators(steps)
steps = _apply_model_params(steps)
base_train = _split_base_samples(spectro)
if allowed_base is not None:
base_train = [sample for sample in base_train if sample in allowed_base]
group_by_sample = _split_group_grain(splitter, spectro, base_train)
if _is_repetition_dataset(spectro):
group_by_sample = group_by_sample or _repetition_grain(spectro, base_train)
# BASE-grain folds: split the base train pool only; train = base-train, val = base-val. The children
# are NEVER listed in a fold (the FoldSet stays a clean base-grain OOF partition); they are pulled
# into fit-train by the host expansion keyed on the origin's fold side. The split runs over the BASE
# rows only (before any child exists), so the fold partition is identical for both augmentation paths.
if base_folds_override is not None:
remaining = set(base_train)
base_folds = [
([sample for sample in train if sample in remaining],
[sample for sample in validation if sample in remaining])
for train, validation in base_folds_override
]
elif _is_repetition_dataset(spectro):
base_folds = _build_group_folds(splitter, spectro, base_train)
else:
base_folds = _build_folds(splitter, spectro, base_train, set())
# GLOBAL stateless augmentation (#8): fit once on the whole train (leakage-free only for stateless
# per-sample augmenters), children shared across all folds (resolver discovers them from identity).
# FOLD-LOCAL augmentation (#32): a STATEFUL/SUPERVISED/balanced augmenter is fit inside EACH fold's
# train only, so each fold (+ the full-train refit) has its OWN children — leakage-safe for the
# stateful case. `fold_children` keys the per-fold expansion; it is pickled for the adapter's resolver.
fold_children: dict[str, dict[int, list[int]]] | None = None
fold_feature_views: dict[str, tuple[Any, dict[int, int], set[int]]] | None = None
if fold_local and (interleaved or late_exclusion_length):
fold_children, augmentation_by_sample_int, fold_feature_views, fold_replay = _build_fold_local_prefix_views(
pipeline[aug_index:after_aug + late_exclusion_length], spectro, base_folds, base_train,
prefix_context, chart_snapshots, chart_transform_snapshots, chart_transform_offset,
)
replay_stages.extend(fold_replay)
elif fold_local:
fold_children, augmentation_by_sample_int = _build_fold_local_children(aug_steps, spectro, base_folds, base_train, prefix_context, chart_snapshots)
elif interleaved and not materialized_early:
replay_stages = _materialize_augmentation_prefix(
pipeline[:after_aug], spectro,
chart_snapshots=chart_snapshots, chart_transform_snapshots=chart_transform_snapshots,
)
elif not materialized_early:
for aug_step in aug_steps:
_apply_sample_augmentation(aug_step, spectro, prefix_context)
if chart_snapshots is not None:
chart_snapshots.append(copy.deepcopy(spectro))
if not fold_local:
_augmented_ints, augmentation_by_sample_int = _augmentation_grain(spectro, "+".join(_augmentation_label(step) for step in aug_steps))
if separation is not None:
branch_step, branch_body = separation
run_dir.mkdir(parents=True, exist_ok=True)
pickle_path = run_dir / "augmented_dataset.pkl"
pickle_path.write_bytes(pickle.dumps({"dataset": spectro, "fold_children": fold_children, "fold_feature_views": fold_feature_views} if fold_local else spectro))
result = _run_separation_branch(
post_aug_steps, branch_step, branch_body, spectro, dataset_arg, cli,
venv_python, run_dir / "separation", metric, task_type,
dataset_pickle=str(pickle_path), config_name=config_name,
random_state=random_state, augmentation_by_sample=augmentation_by_sample_int,
fold_children=fold_children, fold_feature_views=fold_feature_views, folds_override=base_folds,
refit=refit,
)
result = _attach_pre_augmentation_replay(result, replay_stages)
result._dagml_chart_aug_snapshots = chart_snapshots # type: ignore[attr-defined]
result._dagml_chart_transform_snapshots = chart_transform_snapshots # type: ignore[attr-defined]
return result
# Identity is minted on the AUGMENTED dataset so children get their own observation_id + the origin's
# sample_id (augmented=True). The CV universe = base train + the augmented children.
identity = mint_identity(spectro)
samples = [int(s) for s in spectro.index_column("sample", {})]
origins = [int(o) for o in spectro.index_column("origin", {})]
augmented_ints = [sample_int for sample_int, origin_int in zip(samples, origins, strict=True) if sample_int != origin_int]
allowed_origins = set(base_train)
cv_universe = base_train + [sample for sample, origin in zip(samples, origins, strict=True) if sample != origin and origin in allowed_origins]
if group_by_sample:
group_by_sample.update({sample: group_by_sample[origin] for sample, origin in zip(samples, origins, strict=True) if sample != origin and origin in group_by_sample})
envelope = build_envelope(
spectro,
identity,
sample_ints=cv_universe,
augmentation_by_sample=augmentation_by_sample_int,
group_by_sample=group_by_sample,
)
dsl = assemble_cv_refit_dsl(steps, identity, envelope, base_folds, dsl_id="nirs4all-augmentation", n_splits=len(base_folds))
# Operator SELECT and its pruned campaign must derive identical signed
# variants; both runners forward this seed to native selection.
dsl["root_seed"] = random_state if random_state is not None else 0
import dag_ml
graph = dag_ml.compile_pipeline_dsl_artifact_with_controllers(dsl, controller_manifests()).graph.to_dict()
# Operator generators materialize one model node per choice. Bind the
# augmentation envelope to every candidate in the native union graph.
from .cli_runner import data_bindings_for_nodes
model_ids = [node["id"] for node in graph["nodes"] if node["kind"] == "model"]
if len(model_ids) > 1:
dsl["data_bindings"] = data_bindings_for_nodes(model_ids, envelope)
for binding in dsl["data_bindings"][:len(checkpoint_steps)]:
binding["view_policy"] = {"include_augmented_train": False}
# The native FoldSet stays base-grain for OOF validation. Declare the exact
# synthetic observations fitted in each fold separately so only the CV
# in-sample `train` report can include them. A fold-local augmenter creates
# distinct children per fold; a global stateless augmenter shares children
# whose origins belong to the fold's training cohort.
origin_by_sample = dict(zip(samples, origins, strict=True))
augmented_train_by_fold: dict[str, list[str]] = {}
for fold_index, (train_ints, _validation_ints) in enumerate(base_folds):
fold_label = f"fold{fold_index}"
train_set = set(train_ints)
if fold_children is None:
children = [sample for sample in augmented_ints if origin_by_sample[sample] in train_set]
else:
excluded = (fold_feature_views[fold_label][2]
if fold_feature_views is not None and fold_label in fold_feature_views else set())
children = [sample for child_ids in fold_children.get(fold_label, {}).values()
for sample in child_ids if origin_by_sample[sample] in train_set and sample not in excluded]
augmented_train_by_fold[fold_label] = [identity.to_wire(sample) for sample in children]
for binding in dsl["data_bindings"][len(checkpoint_steps):]:
binding.setdefault("view_policy", {}).update({
"include_augmented_cv_train_predictions": True,
"augmented_cv_train_prediction_ids_by_fold": augmented_train_by_fold,
})
run_dir.mkdir(parents=True, exist_ok=True)
pickle_path = run_dir / "augmented_dataset.pkl"
# Fold-local pickles the dataset + the fold→children map (the resolver's per-fold expansion); the
# global path pickles the bare dataset (the resolver discovers dataset-global children from identity).
pickle_path.write_bytes(pickle.dumps({"dataset": spectro, "fold_children": fold_children, "fold_feature_views": fold_feature_views} if fold_local else spectro))
outcome = run_cv_refit_bundle(
dsl=dsl, envelope=envelope, graph=graph, dataset_path=dataset_arg, workdir=run_dir, dagml_cli=cli, venv_python=venv_python, selection_metric=metric, dataset_pickle=str(pickle_path), dataset=spectro, fold_children=fold_children, fold_feature_views=fold_feature_views, random_state=random_state, refit=refit
)
if outcome["returncode"] != 0:
_raise_run_failure(outcome, "dag-ml augmentation run failed")
if checkpoint_steps:
from .result import _variant_cv_score
producer_scores = [
{**outcome["scores"], "reports": [report for report in outcome["scores"]["reports"] if report["producer_node"] == model_id]}
for model_id in model_ids
]
decision = dag_ml.select_candidate(
{"id": "select:augmentation_checkpoints", "metric": {"name": metric, "objective": "maximize" if is_higher_better(metric) else "minimize"}},
[{"candidate_id": str(index), "metrics": {metric: _variant_cv_score(scores, metric)}} for index, scores in enumerate(producer_scores)],
)
result = _project_operator_sweep(
list(zip(producer_scores, [_model_name([step]) for step in steps if isinstance(step, dict) and "model" in step], strict=True)),
spectro.name, metric, task_type, task_type != "regression", [config_name] * len(model_ids),
results_by_index=[
[frame for frame in outcome["results"] if any(
block.get("producer_node") == model_id
for block in (frame.get("result", frame).get("predictions") or [])
)]
for model_id in model_ids
],
identities_by_index=[identity] * len(model_ids),
refit_artifacts_by_index=[
[artifact for artifact in outcome["refit_artifacts"] if f":{model_id}:" in artifact["artifact_id"]]
for model_id in model_ids
],
selected_index=int(decision["selected_candidate_id"]),
emit_all_refits=True,
)
else:
result = _scores_to_run_result(
outcome["scores"],
spectro.name,
_model_name(steps),
metric,
task_type,
config_name=config_name,
results=outcome["results"],
identity=identity,
refit_artifacts=outcome["refit_artifacts"],
)
result = _attach_pre_augmentation_replay(result, replay_stages)
result._dagml_chart_aug_snapshots = chart_snapshots # type: ignore[attr-defined]
result._dagml_chart_transform_snapshots = chart_transform_snapshots # type: ignore[attr-defined]
if capture is not None:
capture.update(
scores=outcome["scores"], results=outcome["results"], identity=identity,
refit_artifacts=result._dagml_refit_artifacts, model_name=_model_name(steps),
)
return result
_MERGE_NODE_ID = "merge:concat"
def _run_separation_branch(pipeline: list[Any], branch_step: dict[str, Any], branch_body: list[Any], spectro: Any, dataset_arg: str, cli: str, venv_python: str, run_dir: Path, metric: str, task_type: str, dataset_pickle: str | None = None, config_name: str = "", random_state: int | None = None, augmentation_by_sample: dict[int, str] | None = None, fold_children: dict[str, dict[int, list[int]]] | None = None, fold_feature_views: dict[str, tuple[Any, dict[int, int], set[int]]] | None = None, folds_override: list[tuple[list[int], list[int]]] | None = None, cv_pool: list[int] | None = None, excluded_sample_ints: set[int] | None = None, refit: bool = True) -> RunResult:
"""Run a by_metadata/by_tag separation branch + concat merge as ONE native dag-ml fan-out run.
Lowers the branch to an ``auto_separate`` template (one branch carrying the criterion + the
model sub-pipeline) followed by a concat ``PredictionJoin`` merge, builds the envelope with the
criterion column's per-sample metadata so the **native** ``fan_out_data_aware_branches`` discovers
the partition values, compiles the fanned graph (one model node per value), and drives dag-ml-cli.
The native concat-merge handler reassembles the per-partition OOF into one full-universe OOF
attributed to the merge node — its cross-fold average is ``cv_best_score``.
The criterion column's metadata is emitted onto the relations; the adapter honors each fanned
model's ``branch_view`` selector (via the sample→metadata map) so each model fits/predicts only
its partition. The legacy nirs4all concat-merge reassembly is broken here (MERGE-E003), so this
native path is a correction — the parity baseline is a direct sklearn-per-partition OOF.
"""
import dag_ml
from nirs4all.pipeline.dagml.cli_runner import data_bindings_for_nodes, split_invocation_for
criterion = branch_step["branch"]
mode, key = ("by_metadata", criterion["by_metadata"]) if "by_metadata" in criterion else ("by_tag", criterion["by_tag"])
# The splitter lives at the top level (before the branch); the branch body is the model
# sub-pipeline applied per partition. Drop the splitter from the body if it slipped in.
_, splitter = _split_pipeline(pipeline)
if splitter is None:
raise DagMlUnsupported("engine='dag-ml' requires a cross-validator step (e.g. KFold) in the pipeline")
# Drop the splitter / None no-ops and reject wavelength-requiring or non-routable ops in the body.
body_steps = _supported_body_steps([step for step in branch_body if not _is_split_step(step)])
identity = mint_identity(spectro)
pool = cv_pool if cv_pool is not None else (_split_base_samples(spectro) if augmentation_by_sample is not None else spectro.index_column("sample", {"partition": "train"}))
excluded = excluded_sample_ints or set()
folds = folds_override if folds_override is not None else (_build_group_folds(splitter, spectro, pool) if _is_repetition_dataset(spectro) else _build_folds(splitter, spectro, pool, excluded))
if isinstance(excluded, FoldLocalExclusion) and not isinstance(folds, FoldLocalFolds):
folds = excluded.apply(spectro, folds)
# Per-sample criterion values: the first map seeds the envelope relations (native fan-out reads
# partition values from it); the second is the adapter's sample_id→metadata map for branch_view.
metadata_by_sample, sample_metadata = _branch_metadata(spectro, identity, mode, key)
group_by_sample = _split_group_grain(splitter, spectro, pool)
if _is_repetition_dataset(spectro):
group_by_sample = group_by_sample or _repetition_grain(spectro, pool)
if augmentation_by_sample is not None:
samples = [int(sample) for sample in spectro.index_column("sample", {})]
origins = [int(origin) for origin in spectro.index_column("origin", {})]
allowed_origins = set(pool)
children = [sample for sample, origin in zip(samples, origins, strict=True) if sample != origin and origin in allowed_origins]
if group_by_sample:
group_by_sample.update({sample: group_by_sample[origin] for sample, origin in zip(samples, origins, strict=True) if sample != origin and origin in group_by_sample})
universe = [*pool, *children]
else:
universe = pool
envelope = build_envelope(spectro, identity, sample_ints=universe, excluded_sample_ints=excluded, metadata_by_sample=metadata_by_sample, augmentation_by_sample=augmentation_by_sample, group_by_sample=group_by_sample)
# Compat auto_separate template: ONE branch (the model sub-pipeline) carrying the criterion +
# mode, marked auto_separate; the native fan-out expands it into N per-partition branches.
template = {"id": "per_partition", "steps": [_branch_compat_step(step) for step in body_steps]}
# Always by_metadata mode: the criterion (whether nirs4all by_metadata or by_tag) is emitted as a
# metadata column on the relations, so the native fan-out discovers its values from there.
has_concat_merge = any(isinstance(step, dict) and step.get("merge") == "concat" for step in pipeline)
compat_dsl: dict[str, Any] = {
"id": "nirs4all-separation-branch",
"pipeline": [
{"branch": {"branches": [template]}, "mode": "by_metadata", "selector": {"metadata_key": key}, "metadata": {"auto_separate": True}},
*([{"merge": "concat", "output_as": "predictions", "id": _MERGE_NODE_ID}] if has_concat_merge else []),
],
}
# NATIVE fan-out (no Python suffix replication): dag-ml reads the partition values from the
# envelope relations and expands one branch per sorted value, owning the node-id suffixing.
compat_dsl["root_seed"] = random_state if random_state is not None else 0
fanned_dsl = dag_ml.fan_out_data_aware_branches(compat_dsl, envelope).to_dict()
manifests = controller_manifests()
fanned_dsl["root_seed"] = random_state if random_state is not None else 0
graph = dag_ml.compile_pipeline_dsl_artifact_with_controllers(fanned_dsl, manifests).graph.to_dict()
model_ids = [node["id"] for node in graph["nodes"] if node["kind"] == "model"]
if not model_ids:
raise DagMlUnsupported("separation-branch fan-out produced no per-partition model nodes")
# Per-partition data_bindings (one per fanned model node) + the materialized fold set. The CLI's
# own fan-out is a no-op on the already-fanned DSL (the auto_separate marker is consumed).
fanned_dsl["data_bindings"] = data_bindings_for_nodes(model_ids, envelope)
fanned_dsl["split_invocation"] = split_invocation_for(identity, folds, n_splits=len(folds))
outcome = run_cv_refit_bundle(
dsl=fanned_dsl, envelope=envelope, graph=graph, dataset_path=dataset_arg, workdir=run_dir, dagml_cli=cli, venv_python=venv_python, selection_metric=metric, sample_metadata=sample_metadata, dataset_pickle=dataset_pickle, dataset=spectro, fold_children=fold_children, fold_feature_views=fold_feature_views, random_state=random_state, refit=refit
)
if outcome["returncode"] != 0:
_raise_run_failure(outcome, "dag-ml separation-branch run failed")
winner_variant_id = next(
(
report.get("variant_id")
for report in (outcome["scores"] or {}).get("reports", [])
if report.get("partition") == "final" and report.get("fold_id") is None
),
None,
)
if winner_variant_id is None and not refit:
winner_variant_id = next(
(report.get("variant_id") for report in (outcome["scores"] or {}).get("reports", [])
if report.get("partition") == "validation" and report.get("fold_id") != "avg"),
None,
)
results_by_variant = _frames_by_variant(outcome["results"], winner_variant_id) if winner_variant_id is not None else None
# The concat-merge producer's reports carry both the full-universe cross-fold OOF average
# (`cv_best_score`) AND a reassembled `(test, fold_id=None)` block (`best_rmse`): dag-ml's native
# off-fold merge handler reassembles each per-partition refit model's held-out TEST prediction
# (the node runner emits it with `fold_id=None`) into one full-universe test block under the merge
# node. Both scores are the separation branch's, surfaced by `_scores_to_run_result`.
if has_concat_merge:
result = _scores_to_run_result(outcome["scores"], spectro.name, _model_name(body_steps), metric, task_type, producer=_MERGE_NODE_ID, config_name=config_name, refit_artifacts=outcome["refit_artifacts"])
else:
predictions = Predictions()
for branch_index, model_id in enumerate(model_ids):
local = _scores_to_run_result(
outcome["scores"], spectro.name, _model_name(body_steps), metric, task_type,
producer=model_id, config_name=config_name,
results_by_variant=results_by_variant, identity=identity,
)
for row in local.predictions.filter_predictions(load_arrays=True):
row["branch_id"] = branch_index
row["branch_name"] = str(((next(node for node in graph["nodes"] if node["id"] == model_id).get("metadata") or {}).get("dsl_branch_selector") or {}).get("metadata", {}).get(key, branch_index))
predictions.extend_from_list([row])
predictions.flush()
result = RunResult(predictions=predictions, per_dataset={spectro.name: {"engine": "dag-ml"}})
result._dagml_score_set = outcome["scores"] # noqa: SLF001
result._dagml_refit_artifacts = outcome["refit_artifacts"] # noqa: SLF001
from .native_results import separation_replay_manifest
replay_manifest = separation_replay_manifest(graph, key, outcome["refit_artifacts"])
if replay_manifest is not None:
result.per_dataset[spectro.name]["separation_replay"] = replay_manifest
return result
def _branch_compat_step(step: Any) -> dict[str, Any]:
"""Lower one branch-body step (transform / {"model": M} / {"y_processing": op}) to compat DSL."""
from nirs4all.pipeline.dagml_bridge import _step_to_dsl
return _step_to_dsl(step)
def _branch_metadata(spectro: Any, identity: Any, mode: str, key: str) -> tuple[dict[str, dict[int, Any]], dict[str, dict[str, Any]]]:
"""Build the criterion's per-sample metadata: ``({col: {sample_int: value}}, {wire_id: {col: value}})``.
The first map seeds the envelope relations (native fan-out reads partition values from it); the
second is the adapter's ``sample_id → metadata`` map for honoring each branch's ``branch_view``.
A ``by_tag`` criterion is represented as a metadata column (the tag value per sample), matching
the envelope's metadata-carried relations.
"""
sample_ints = [int(value) for value in spectro.index_column("sample", {})]
origin_ints = [int(value) for value in spectro.index_column("origin", {})]
values = spectro.metadata_column(key, {}) if mode == "by_metadata" else spectro.get_tag(key, {})
base_ints = [sample for sample, origin in zip(sample_ints, origin_ints, strict=True) if sample == origin]
if len(values) == len(sample_ints):
by_int = {sample_int: (str(value) if value is not None else None) for sample_int, value in zip(sample_ints, values, strict=True)}
elif len(values) == len(base_ints):
by_int = {sample_int: (str(value) if value is not None else None) for sample_int, value in zip(base_ints, values, strict=True)}
else:
raise ValueError(f"branch metadata {key!r} has {len(values)} values for {len(base_ints)} base rows")
for sample_int, origin_int in zip(sample_ints, origin_ints, strict=True):
if sample_int != origin_int and by_int.get(sample_int) is None:
by_int[sample_int] = by_int[origin_int]
metadata_by_sample = {key: dict(by_int)}
sample_metadata = {identity.to_wire(sample_int): {key: value} for sample_int, value in by_int.items()}
return metadata_by_sample, sample_metadata
_FUSION_MERGE_NODE_ID = "merge:fusion"
def _canonical_branch_step(step: Any, node_id: str) -> dict[str, Any]:
"""Lower one branch-body step to a CANONICAL pipeline-DSL step (``kind`` + ``operator.class``).
The duplication-fusion path emits a *canonical* ``steps`` DSL (not the compat ``pipeline`` form):
dag-ml's nirs4all-compat importer whitelists merge modes (``concat``/``predictions``/… only) and
REFUSES ``fusion`` — but the canonical ``PipelineDslMergeStep.merge_mode`` is a free-form string the
runtime reads verbatim, so the fusion merge must be expressed canonically. The branch bodies must
therefore also be canonical. This reuses the verified compat lowering
(:func:`~nirs4all.pipeline.dagml_bridge._step_to_dsl` — operator FQN + JSON-safe params +
native param-generator entries) and re-keys it into the canonical ``{"kind", "id", "operator":
{"class": …}, "params": …}`` shape per node kind (transform / y_transform / model), assigning the
explicit ``node_id`` so each branch's nodes are uniquely named.
"""
from nirs4all.pipeline.dagml_bridge import _step_to_dsl
compat = _step_to_dsl(step)
if "model" in compat:
out: dict[str, Any] = {"kind": "model", "id": node_id, "operator": {"class": compat["model"]}, "params": compat.get("params", {})}
if "generators" in compat:
out["generators"] = compat["generators"]
if "metadata" in compat:
out["metadata"] = compat["metadata"]
return out
if "y_processing" in compat:
inner = compat["y_processing"]
return {"kind": "y_transform", "id": node_id, "operator": {"class": inner["class"]}, "params": inner.get("params", {})}
if "preprocessing" in compat:
inner = compat["preprocessing"]
return {
"kind": "transform", "id": node_id,
"operator": {"class": inner["class"]}, "params": inner.get("params", {}),
**({"metadata": compat["metadata"]} if "metadata" in compat else {}),
**({"shape": compat["shape"]} if "shape" in compat else {}),
}
# Bare transform: compat is {"class": FQN, "params": {...}}.
return {"kind": "transform", "id": node_id, "operator": {"class": compat["class"]}, "params": compat.get("params", {})}
def _canonical_branch(branch_body: list[Any], branch_index: int) -> dict[str, Any]:
"""Lower one duplication sub-pipeline (a list of steps) to a canonical branch with unique node ids.
``None`` no-ops are dropped (never lowered to a ``builtins.NoneType`` node) and wavelength-requiring /
non-routable ops in the body raise a catchable :class:`DagMlUnsupported` — same coverage guarantee
the top-level path gives, applied here so duplication / by_source / stacking branch bodies are safe.
"""
steps = _supported_body_steps([step for step in branch_body if not _is_split_step(step)])
return {
"id": f"branch_{branch_index}",
"steps": [_canonical_branch_step(step, f"branch:{branch_index}.node:{node_index}") for node_index, step in enumerate(steps)],
}
def _run_checkpoint_before_duplication_branch(
pipeline: list[Any], branches: list[list[Any]], first_model: dict[str, Any], last_model: dict[str, Any],
spectro: Any, dataset_arg: str, cli: str, venv_python: str, run_dir: Path,
metric: str, task_type: str, dataset_pickle: str | None, config_name: str,
random_state: int | None, refit: bool,
) -> RunResult:
"""Run the prior checkpoint and each branch model in one native DAG campaign."""
import dag_ml
from .cli_runner import data_bindings_for_nodes, split_invocation_for
_, splitter = _split_pipeline(pipeline)
identity = mint_identity(spectro)
pool = list(spectro.index_column("sample", {"partition": "train"}))
folds = _build_folds(splitter, spectro, pool, set())
envelope = build_envelope(spectro, identity, sample_ints=pool,
group_by_sample=_split_group_grain(splitter, spectro, pool))
prior_id = "model:checkpoint_before_branch"
canonical_dsl: dict[str, Any] = {
"id": "nirs4all-checkpoint-duplication",
"steps": [
_canonical_branch_step(first_model, prior_id),
{"kind": "branch", "mode": "duplication", "branches": [
_canonical_branch([*body, last_model], index) for index, body in enumerate(branches)
]},
],
}
canonical_dsl["root_seed"] = random_state if random_state is not None else 0
graph = dag_ml.compile_pipeline_dsl_artifact_with_controllers(canonical_dsl, controller_manifests()).graph.to_dict()
producer_ids = [node["id"] for node in graph["nodes"] if node["kind"] == "model"]
if len(producer_ids) != len(branches) + 1 or producer_ids[0] != prior_id:
raise DagMlUnsupported("checkpoint branch compilation did not preserve every model producer")
canonical_dsl["data_bindings"] = data_bindings_for_nodes(producer_ids, envelope)
canonical_dsl["split_invocation"] = split_invocation_for(identity, folds, n_splits=len(folds))
outcome = run_cv_refit_bundle(
dsl=canonical_dsl, envelope=envelope, graph=graph, dataset_path=dataset_arg,
workdir=run_dir, dagml_cli=cli, venv_python=venv_python,
selection_metric=metric, dataset_pickle=dataset_pickle, dataset=spectro,
random_state=random_state, refit=refit,
)
if outcome["returncode"] != 0:
_raise_run_failure(outcome, "dag-ml checkpoint branch run failed")
predictions = Predictions()
for index, producer in enumerate(producer_ids):
model_name = _model_name([first_model if index == 0 else last_model])
frames = [frame for frame in outcome["results"]
if (frame.get("result") if frame.get("type") == "result" else frame).get("node_id") == producer]
projected = _scores_to_run_result(
outcome["scores"], spectro.name, model_name, metric, task_type,
producer=producer, config_name=config_name, results=frames,
identity=identity, refit_artifacts=outcome["refit_artifacts"],
)
for row in projected.predictions.filter_predictions(load_arrays=True):
if index:
row["branch_id"] = index - 1
row["branch_name"] = f"branch_{index - 1}"
predictions.extend_from_list([row])
predictions.flush()
result = RunResult(predictions=predictions, per_dataset={spectro.name: {
"engine": "dag-ml", "refit_enabled": refit,
"checkpoint_producers": producer_ids,
}})
result._dagml_score_set = outcome["scores"] # noqa: SLF001 - preserve native multi-producer evidence
result._dagml_refit_artifacts = outcome["refit_artifacts"] # noqa: SLF001
result._dagml_node_results = outcome["results"] # noqa: SLF001
return result
def _run_checkpoint_inside_duplication_feature_merge(
pipeline: list[Any], branches: list[list[Any]], first_model: dict[str, Any], last_model: dict[str, Any],
spectro: Any, dataset_arg: str, cli: str, venv_python: str, run_dir: Path,
metric: str, task_type: str, dataset_pickle: str | None, config_name: str,
random_state: int | None, refit: bool,
) -> RunResult:
"""Keep branch-local model scores and fitted feature joins in one native graph."""
import dag_ml
from .cli_runner import data_bindings_for_nodes, split_invocation_for
_, splitter = _split_pipeline(pipeline)
identity = mint_identity(spectro)
pool = list(spectro.index_column("sample", {"partition": "train"}))
folds = _build_folds(splitter, spectro, pool, set())
envelope = build_envelope(spectro, identity, sample_ints=pool,
group_by_sample=_split_group_grain(splitter, spectro, pool))
canonical_dsl: dict[str, Any] = {
"id": "nirs4all-checkpoint-duplication-feature-merge",
"steps": [
{"kind": "branch", "mode": "duplication", "branches": [
_canonical_branch([*body, first_model, last_model], index) for index, body in enumerate(branches)
]},
{"kind": "merge", "id": "merge:features", "merge_mode": "concat",
"output_as": "features", "include_original_data": False},
],
}
canonical_dsl["root_seed"] = random_state if random_state is not None else 0
graph = dag_ml.compile_pipeline_dsl_artifact_with_controllers(canonical_dsl, controller_manifests()).graph.to_dict()
producer_ids = [node["id"] for node in graph["nodes"] if node["kind"] == "model"]
if len(producer_ids) != 2 * len(branches):
raise DagMlUnsupported("checkpoint feature merge compilation did not preserve every model producer")
bound_ids = [node["id"] for node in graph["nodes"] if node["kind"] in {"transform", "model"}]
canonical_dsl["data_bindings"] = data_bindings_for_nodes(bound_ids, envelope)
canonical_dsl["split_invocation"] = split_invocation_for(identity, folds, n_splits=len(folds))
outcome = run_cv_refit_bundle(
dsl=canonical_dsl, envelope=envelope, graph=graph, dataset_path=dataset_arg,
workdir=run_dir, dagml_cli=cli, venv_python=venv_python,
selection_metric=metric, dataset_pickle=dataset_pickle, dataset=spectro,
random_state=random_state, refit=refit,
)
if outcome["returncode"] != 0:
_raise_run_failure(outcome, "dag-ml checkpoint feature merge run failed")
predictions = Predictions()
for index, producer in enumerate(producer_ids):
model_name = _model_name([first_model if index % 2 == 0 else last_model])
frames = [frame for frame in outcome["results"]
if (frame.get("result") if frame.get("type") == "result" else frame).get("node_id") == producer]
projected = _scores_to_run_result(
outcome["scores"], spectro.name, model_name, metric, task_type,
producer=producer, config_name=config_name, results=frames,
identity=identity, refit_artifacts=outcome["refit_artifacts"],
)
for row in projected.predictions.filter_predictions(load_arrays=True):
branch_index = index // 2
if index % 2 == 0:
if row["fold_id"] != "final":
row["branch_id"] = branch_index
row["branch_name"] = f"branch_{branch_index}"
elif branch_index and row["fold_id"] != "final":
# Legacy computes the second Ridge only in its terminal branch replay.
continue
predictions.extend_from_list([row])
predictions.flush()
result = RunResult(predictions=predictions, per_dataset={spectro.name: {
"engine": "dag-ml", "refit_enabled": refit, "checkpoint_producers": producer_ids,
}})
result._dagml_score_set = outcome["scores"] # noqa: SLF001
result._dagml_refit_artifacts = outcome["refit_artifacts"] # noqa: SLF001
result._dagml_node_results = outcome["results"] # noqa: SLF001
return result
def _branch_merge_transformer_step(branches: list[list[Any]], merge_mode: str) -> DuplicationBranchMergeTransformer:
"""Build the importable transformer step used for native duplication feature/all merges."""
from nirs4all.pipeline.dagml_bridge import _json_safe_params, _qualname
if merge_mode not in ("features", "all"):
raise DagMlUnsupported("engine='dag-ml' supports duplication branch merge only for merge='features' or merge='all' in this slice")
lowered_branches: list[dict[str, Any]] = []
for branch in branches:
lowered_steps: list[dict[str, Any]] = []
for step in _supported_body_steps([part for part in branch if not _is_split_step(part)]):
if isinstance(step, dict) and "model" in step:
if merge_mode != "all":
raise DagMlUnsupported(
"engine='dag-ml' supports duplication merge=features only for feature-only branch bodies; "
"branch-local models require merge='all'."
)
model = step["model"]
lowered_steps.append({"kind": "model", "class": _qualname(model), "params": _json_safe_params(model)})
elif isinstance(step, dict):
raise DagMlUnsupported(
"engine='dag-ml' supports duplication feature/all merges only for plain branch transform steps "
"and merge='all' branch-local model steps; "
"structured branch step dictionaries "
f"such as {sorted(step)} stay on the legacy path."
)
else:
lowered_steps.append({"kind": "transform", "class": _qualname(step), "params": _json_safe_params(step)})
lowered_branches.append({"steps": lowered_steps})
return DuplicationBranchMergeTransformer(branches=lowered_branches, merge_mode=merge_mode)
def _branch_fusion_model_step(branches: list[list[Any]], aggregate: str, task_type: str) -> dict[str, Any]:
"""Represent legacy mean-fusion branches as one cloneable native model task."""
if aggregate != "mean" or task_type != "regression":
raise DagMlUnsupported("augmentation with branch model fusion currently supports regression mean only")
lowered = _branch_merge_transformer_step(branches, "all")
return {"model": DuplicationFusionEstimator(lowered.branches)}
def _canonical_source_branch(branch_body: list[Any], source_index: int, *, force_generator: bool = False) -> dict[str, Any]:
"""Lower the shared by_source body to a canonical branch BOUND to one source (S4).
Same lowering as :func:`_canonical_branch` (the shared model sub-pipeline, unique node ids per
source), but every MODEL node carries ``metadata.source_index`` so the node runner materializes
ONLY that source's feature block — late fusion by source. The branch index IS the source index
(one branch per source), so a fold view stays full-sample (all branches see all samples) while
each branch's model sees a different source's columns.
"""
steps = _supported_body_steps([step for step in branch_body if not _is_split_step(step)])
branch: dict[str, Any]
generator_indices = [index for index, step in enumerate(steps) if isinstance(step, dict) and "_or_" in step]
if generator_indices:
if len(generator_indices) != 1:
raise DagMlUnsupported("by_source supports one _or_ operator generator per source branch")
generator_index = generator_indices[0]
generator = steps[generator_index]
if set(generator) != {"_or_"} or not isinstance(generator["_or_"], list) or not generator["_or_"]:
raise DagMlUnsupported("by_source operator generation requires a bare nonempty _or_ list")
choice_branches: list[dict[str, Any]] = [
{"id": f"choice_{choice_index}", "steps": [_canonical_branch_step(choice, f"source:{source_index}.choice:{choice_index}")]}
for choice_index, choice in enumerate(generator["_or_"])
]
if any(choice["steps"][0]["kind"] != "transform" for choice in choice_branches):
raise DagMlUnsupported("by_source _or_ choices must be X transforms")
tail = [
_canonical_branch_step(step, f"source:{source_index}.tail:{index}")
for index, step in enumerate(steps[generator_index + 1:])
]
if not tail or tail[-1]["kind"] != "model":
raise DagMlUnsupported("by_source _or_ must be followed by a branch-local model")
prefix = [
_canonical_branch_step(step, f"source:{source_index}.prefix:{index}")
for index, step in enumerate(steps[:generator_index])
]
branch = {"id": f"branch_{source_index}", "steps": [*prefix, {
"kind": "generator", "id": f"generator:source_{source_index}", "mode": "or",
"branches": choice_branches, "tail": tail,
}]}
else:
branch = _canonical_branch(steps, source_index)
if force_generator:
# A one-choice generator makes the static source an explicit member
# of the same native Cartesian SELECT as the generated source.
branch["steps"] = [{
"kind": "generator", "id": f"generator:source_{source_index}",
"mode": "or", "branches": [{"id": "only", "steps": branch["steps"]}],
}]
def bind_source(nodes: list[dict[str, Any]]) -> None:
for node in nodes:
if node["kind"] == "model":
node["metadata"] = {**node.get("metadata", {}), "source_index": source_index}
elif node["kind"] == "generator":
bind_source(node.get("tail", []))
for choice in node["branches"]:
bind_source(choice["steps"])
bind_source(branch["steps"])
return branch
def _source_preprocessing_metadata(source_steps: dict[str, list[Any]], source_layout: dict[str, Any] | None) -> dict[str, Any]:
"""Lower per-source preprocessing dict by explicit ``source_layout.source_order``.
The legacy by_source dict keys are user-facing source names (``source_0`` unless the
dataset exposes names). The native data plan separately names the materialized blocks
``src0``/``src1``. This helper consumes the explicit layout and rejects any mismatch
instead of inferring an order from the dict body.
"""
if not isinstance(source_layout, dict):
raise DagMlUnsupported("by_source distinct preprocessing requires envelope plan.source_layout")
source_order = source_layout.get("source_order")
source_ids = source_layout.get("source_ids")
if not (isinstance(source_order, list) and all(isinstance(item, str) for item in source_order)):
raise DagMlUnsupported("by_source distinct preprocessing requires source_layout.source_order")
if not (isinstance(source_ids, list) and len(source_ids) == len(source_order) and all(isinstance(item, str) for item in source_ids)):
raise DagMlUnsupported("by_source distinct preprocessing requires source_layout.source_ids aligned to source_order")
if len(set(source_order)) != len(source_order):
raise DagMlUnsupported("by_source distinct preprocessing source_layout.source_order contains duplicate names")
if set(source_steps) != set(source_order):
raise DagMlUnsupported(
"by_source distinct preprocessing keys must exactly match source_layout.source_order: "
f"expected={source_order!r} actual={list(source_steps)!r}"
)
sources: list[dict[str, Any]] = []
for source_index, (source_name, source_id) in enumerate(zip(source_order, source_ids, strict=True)):
body_steps = _supported_body_steps(source_steps[source_name])
lowered = [_canonical_branch_step(step, f"source:{source_index}.pre:{step_index}") for step_index, step in enumerate(body_steps)]
if any(step["kind"] != "transform" for step in lowered):
raise DagMlUnsupported("by_source distinct preprocessing bodies may contain only X transforms")
sources.append(
{
"source_name": source_name,
"source_id": source_id,
"source_index": source_index,
"steps": lowered,
}
)
return {
"mode": "by_source_distinct_preproc_concat",
# Legacy by_source merge writes the concatenated block into source 0 but
# leaves non-primary sources present for the downstream multi-source fit.
"preserve_legacy_sources_after_merge": True,
"sources": sources,
}
_PREDICTION_CLONE_FIELDS = (
"dataset_name",
"dataset_path",
"config_name",
"config_path",
"pipeline_uid",
"step_idx",
"op_counter",
"model_name",
"model_classname",
"model_path",
"fold_id",
"sample_indices",
"weights",
"metadata",
"partition",
"y_true",
"y_pred",
"y_proba",
"val_score",
"test_score",
"train_score",
"metric",
"task_type",
"n_samples",
"n_features",
"preprocessings",
"best_params",
"scores",
"branch_id",
"branch_path",
"branch_name",
"exclusion_count",
"exclusion_rate",
"model_artifact_id",
"trace_id",
"refit_context",
"target_processing",
"result_metadata",
)
def _repeat_by_source_merge_projection(result: RunResult, n_sources: int) -> RunResult:
"""Match legacy by_source+concat bookkeeping: one identical result block per source."""
if n_sources <= 1:
return result
rows = result.predictions.filter_predictions(load_arrays=True)
cv_rows = [row for row in rows if str(row.get("fold_id")) != "final"]
final_rows = [row for row in rows if str(row.get("fold_id")) == "final"]
repeated = Predictions()
def clone(row: dict[str, Any]) -> None:
payload: dict[str, Any] = {field: row.get(field) for field in _PREDICTION_CLONE_FIELDS if field in row}
repeated.add_prediction(**payload)
for row_group in (cv_rows, final_rows):
for _source_index in range(n_sources):
for row in row_group:
clone(row)
repeated.flush()
result.predictions = repeated
return result
def _duplication_branch_names(pipeline: list[Any], count: int) -> list[str]:
"""Recover legacy branch labels for named duplication branch rows."""
config_keys = {"parallel", "n_jobs", "mode", "strategy", "merge", "output_as", "aggregate", "branch_type", "preserve_branch_context"}
branch_spec = next((step.get("branch") for step in pipeline if isinstance(step, dict) and "branch" in step), None)
if isinstance(branch_spec, dict) and not ({"by_metadata", "by_tag", "by_source"} & set(branch_spec)):
names = [str(name) for name in branch_spec if name not in config_keys and not str(name).startswith("_")]
if len(names) == count:
return names
return [f"branch_{index}" for index in range(count)]
def _prediction_rows_without_refit(result: RunResult) -> list[dict[str, Any]]:
"""Copy prediction rows, excluding standalone refit rows legacy does not emit for merge='all'."""
rows: list[dict[str, Any]] = []
for row in result.predictions.filter_predictions(load_arrays=True):
if str(row.get("fold_id")) == "final":
continue
rows.append(dict(row))
return rows
def _combine_duplication_merge_all_rows(
branch_results: list[RunResult],
branch_names: list[str],
downstream_result: RunResult,
downstream_model_name: str,
dataset_name: str,
) -> RunResult:
"""Compose legacy-shaped rows for duplication ``merge='all'``.
Legacy reports every branch-local model block, then the downstream model trained on
``[branch_features..., branch_predictions...]``. The native transformer supplies the same downstream
design matrix; this helper mirrors the legacy bookkeeping by retaining branch model rows separately.
"""
combined_rows: list[dict[str, Any]] = []
for branch_index, (branch_result, branch_name) in enumerate(zip(branch_results, branch_names, strict=True)):
for row in _prediction_rows_without_refit(branch_result):
row["branch_id"] = branch_index
row["branch_name"] = branch_name
combined_rows.append(row)
for row in _prediction_rows_without_refit(downstream_result):
row["model_name"] = downstream_model_name
row["model_classname"] = downstream_model_name
row["branch_id"] = None
row["branch_name"] = ""
combined_rows.append(row)
predictions = Predictions()
predictions.extend_from_list(combined_rows)
predictions.flush()
return RunResult(predictions=predictions, per_dataset={dataset_name: {"engine": "dag-ml"}})
def _clone_operator_instance(operator: Any) -> Any:
"""Clone sklearn-style operators, falling back to deepcopy for local estimators."""
import copy
from sklearn.base import clone
try:
return clone(operator)
except Exception: # noqa: BLE001 - some local operators are clone-hostile but deepcopy-safe.
return copy.deepcopy(operator)
def _dataset_y_rows(spectro: Any, sample_ints: list[int]) -> np.ndarray:
"""Target rows in the requested sample-int order."""
block = np.asarray(spectro.y({"sample": sample_ints}, include_augmented=False), dtype=float)
stored = spectro.index_column("sample", {"sample": sample_ints})
row_of = {int(sample_int): row for row, sample_int in enumerate(stored)}
return block[[row_of[int(sample_int)] for sample_int in sample_ints]]
def _score_block(y_true: np.ndarray, y_pred: np.ndarray, task_type: str) -> dict[str, float]:
metric_task = "regression" if task_type == "regression" else task_type
metric_names = get_default_metrics(metric_task)
values = eval_list(y_true, y_pred, metric_names)
return {name: float(value) for name, value in zip(metric_names, values, strict=False) if value is not None}
def _transform_branch_matrix(spectro: Any, sample_ints: list[int], transforms: list[Any]) -> np.ndarray:
current = np.asarray(spectro.x_rows(sample_ints, layout="2d"))
for transform in transforms:
current = np.asarray(transform.transform(current))
return current
def _fit_global_branch_transforms(branch: list[Any], spectro: Any, train_pool: list[int]) -> tuple[list[Any], Any, np.ndarray]:
"""Fit branch transforms on full train, matching legacy branch-row bookkeeping for merge='all'."""
current = np.asarray(spectro.x_rows(train_pool, layout="2d"))
y_train = _dataset_y_rows(spectro, train_pool)
transforms: list[Any] = []
model = None
for step in _supported_body_steps([part for part in branch if not _is_split_step(part)]):
if isinstance(step, dict) and "model" in step:
model = step["model"]
break
if isinstance(step, dict):
raise DagMlUnsupported(f"unsupported duplication merge='all' branch step {sorted(step)}")
transform = _clone_operator_instance(step)
transform.fit(current, y_train)
current = np.asarray(transform.transform(current))
transforms.append(transform)
if model is None:
raise DagMlUnsupported("duplication merge='all' requires every branch to contain a branch-local model")
return transforms, model, current
def _add_scored_prediction_rows(
predictions: Predictions,
*,
spectro: Any,
config_name: str,
step_idx: int = 0,
op_counter: int = 0,
model_name: str,
model_classname: str,
fold_id: str,
scores: dict[str, dict[str, float]],
y_by_partition: dict[str, np.ndarray],
pred_by_partition: dict[str, np.ndarray],
samples_by_partition: dict[str, list[int]],
metric: str,
task_type: str,
n_features: int,
branch_id: int | None,
branch_name: str,
weights: np.ndarray | None = None,
) -> None:
score_provenance = {
partition: {
"partition": partition,
"fold_id": fold_id,
"purpose": "measurement",
}
for partition in scores
}
for partition in ("train", "val", "test"):
y_true = y_by_partition.get(partition)
y_pred = pred_by_partition.get(partition)
if y_true is None or y_pred is None or len(y_true) == 0 or len(y_pred) == 0:
continue
predictions.add_prediction(
dataset_name=spectro.name,
dataset_path=spectro.name,
config_name=config_name,
config_path=f"{spectro.name}/{config_name or 'config'}",
step_idx=step_idx,
op_counter=op_counter,
model_name=model_name,
model_classname=model_classname,
fold_id=fold_id,
partition=partition,
sample_indices=samples_by_partition.get(partition),
weights=weights,
y_true=y_true,
y_pred=y_pred,
val_score=scores.get("val", {}).get(metric),
test_score=scores.get("test", {}).get(metric),
train_score=scores.get("train", {}).get(metric),
metric=metric,
task_type=task_type,
n_samples=len(y_true),
n_features=n_features,
scores=scores,
result_metadata={"dagml_host_projection": {
"schema": "nirs4all.dagml-host-projection.v1",
"executor": "python",
"score_provenance": score_provenance,
}},
branch_id=branch_id,
branch_name=branch_name,
)
def _run_duplication_merge_all_branch_result(
branch: list[Any],
branch_index: int,
branch_name: str,
splitter: Any,
spectro: Any,
metric: str,
task_type: str,
config_name: str,
) -> RunResult:
"""Project branch rows with preprocessing fitted only on each fold's train."""
train_pool = [int(sample_int) for sample_int in spectro.index_column("sample", {"partition": "train"})]
test_pool = [int(sample_int) for sample_int in spectro.index_column("sample", {"partition": "test"})]
folds = _build_folds(splitter, spectro, train_pool, set())
model_step, X_train, X_test, preprocessing = _prepare_named_branch_feature_matrix(branch, spectro, train_pool, test_pool)
return _run_model_on_precomputed_matrix(
X_train, X_test,
_dataset_y_rows(spectro, train_pool),
_dataset_y_rows(spectro, test_pool) if test_pool else np.empty((0, 1)),
train_pool, test_pool, folds, model_step["model"], spectro,
metric, task_type, config_name,
branch_id=branch_index, branch_name=branch_name,
preprocessing=preprocessing,
)
def _run_model_on_precomputed_matrix(
X_train_all: np.ndarray,
X_test: np.ndarray,
y_train_all: np.ndarray,
y_test: np.ndarray,
train_pool: list[int],
test_pool: list[int],
folds: list[tuple[list[int], list[int]]],
model_template: Any,
spectro: Any,
metric: str,
task_type: str,
config_name: str,
*,
branch_id: int | None,
branch_name: str,
include_refit: bool = False,
preprocessing: list[Any] | None = None,
matrix_for_fit: Callable[[list[int]], tuple[np.ndarray, np.ndarray]] | None = None,
) -> RunResult:
"""Run legacy-shaped CV rows for a model over an already materialized feature matrix."""
train_position = {sample_int: position for position, sample_int in enumerate(train_pool)}
predictions = Predictions()
fold_models: list[Any] = []
fold_val_samples: list[list[int]] = []
fold_val_scores: list[float] = []
cv_refit_score: float | None = None
model_name = type(model_template).__name__
model_template = _with_fold_preprocessing(model_template, preprocessing)
n_features = int(X_train_all.shape[1]) if X_train_all.ndim > 1 else 1
fold_matrices: list[tuple[np.ndarray, np.ndarray]] = []
for fold_index, (fold_train, fold_val) in enumerate(folds):
train_rows = [train_position[int(sample_int)] for sample_int in fold_train]
val_rows = [train_position[int(sample_int)] for sample_int in fold_val]
fold_X, fold_test_X = matrix_for_fit(fold_train) if matrix_for_fit is not None else (X_train_all, X_test)
fold_matrices.append((fold_X, fold_test_X))
model = _clone_operator_instance(model_template)
model.fit(fold_X[train_rows], _dataset_y_rows(spectro, fold_train))
fold_models.append(model)
fold_val_samples.append([int(sample_int) for sample_int in fold_val])
pred_by_partition = {
"train": np.asarray(model.predict(fold_X[train_rows])),
"val": np.asarray(model.predict(fold_X[val_rows])),
"test": np.asarray(model.predict(fold_test_X)) if len(test_pool) else np.array([]),
}
y_by_partition = {
"train": _dataset_y_rows(spectro, fold_train),
"val": _dataset_y_rows(spectro, fold_val),
"test": y_test,
}
scores = {partition: _score_block(y_by_partition[partition], pred_by_partition[partition], task_type) for partition in pred_by_partition if len(y_by_partition[partition]) and len(pred_by_partition[partition])}
fold_val_scores.append(float(scores["val"][metric]))
_add_scored_prediction_rows(
predictions,
spectro=spectro,
config_name=config_name,
model_name=model_name,
model_classname=model_name,
fold_id=str(fold_index),
scores=scores,
y_by_partition=y_by_partition,
pred_by_partition=pred_by_partition,
samples_by_partition={"train": [int(sample_int) for sample_int in fold_train], "val": [int(sample_int) for sample_int in fold_val], "test": test_pool},
metric=metric,
task_type=task_type,
n_features=n_features,
branch_id=branch_id,
branch_name=branch_name,
)
if len(fold_models) > 1:
weights = EnsembleUtils._scores_to_weights(np.asarray(fold_val_scores, dtype=float), higher_is_better=is_higher_better(metric))
all_train_preds = [np.asarray(model.predict(matrices[0])) for model, matrices in zip(fold_models, fold_matrices, strict=True)]
all_test_preds = [np.asarray(model.predict(matrices[1])) for model, matrices in zip(fold_models, fold_matrices, strict=True)] if len(test_pool) else []
oof_val_preds = [
np.asarray(model.predict(matrices[0][[train_position[int(sample_int)] for sample_int in fold_val]]))
for model, fold_val, matrices in zip(fold_models, fold_val_samples, fold_matrices, strict=True)
]
y_val_oof = np.vstack([_dataset_y_rows(spectro, fold_val) for fold_val in fold_val_samples])
val_samples = [sample_int for fold_val in fold_val_samples for sample_int in fold_val]
pred_sets = {
"avg": (
{"train": np.mean(all_train_preds, axis=0), "val": np.concatenate(oof_val_preds) if oof_val_preds else np.array([]), "test": np.mean(all_test_preds, axis=0) if all_test_preds else np.array([])},
None,
),
"w_avg": (
{
"train": np.sum([weight * pred for weight, pred in zip(weights, all_train_preds, strict=False)], axis=0),
"val": np.concatenate(oof_val_preds) if oof_val_preds else np.array([]),
"test": np.sum([weight * pred for weight, pred in zip(weights, all_test_preds, strict=False)], axis=0) if all_test_preds else np.array([]),
},
weights,
),
}
y_by_partition = {"train": y_train_all, "val": y_val_oof, "test": y_test}
samples_by_partition = {"train": train_pool, "val": val_samples, "test": test_pool}
for fold_id, (pred_by_partition, row_weights) in pred_sets.items():
scores = {
partition: _score_block(y_by_partition[partition], pred_by_partition[partition], task_type)
for partition in pred_by_partition
if len(y_by_partition[partition]) and len(pred_by_partition[partition])
}
if fold_id == "avg":
cv_refit_score = scores.get("val", {}).get(metric)
_add_scored_prediction_rows(
predictions,
spectro=spectro,
config_name=config_name,
model_name=model_name,
model_classname=model_name,
fold_id=fold_id,
scores=scores,
y_by_partition=y_by_partition,
pred_by_partition=pred_by_partition,
samples_by_partition=samples_by_partition,
metric=metric,
task_type=task_type,
n_features=n_features,
branch_id=branch_id,
branch_name=branch_name,
weights=row_weights,
)
if include_refit:
if matrix_for_fit is not None:
X_train_all, X_test = matrix_for_fit(train_pool)
model = _clone_operator_instance(model_template)
model.fit(X_train_all, y_train_all)
pred_by_partition = {
"train": np.asarray(model.predict(X_train_all)),
"test": np.asarray(model.predict(X_test)) if len(test_pool) else np.array([]),
}
y_by_partition = {
"train": y_train_all,
"test": y_test,
}
samples_by_partition = {"train": train_pool, "test": test_pool}
scores = {
partition: _score_block(y_by_partition[partition], pred_by_partition[partition], task_type)
for partition in pred_by_partition
if len(y_by_partition[partition]) and len(pred_by_partition[partition])
}
if cv_refit_score is not None:
scores["val"] = {metric: cv_refit_score}
_add_scored_prediction_rows(
predictions,
spectro=spectro,
config_name=f"{config_name}_refit" if config_name else "",
model_name=model_name,
model_classname=model_name,
fold_id="final",
scores=scores,
y_by_partition=y_by_partition,
pred_by_partition=pred_by_partition,
samples_by_partition=samples_by_partition,
metric=metric,
task_type=task_type,
n_features=n_features,
branch_id=branch_id,
branch_name=branch_name,
)
predictions.flush()
return RunResult(predictions=predictions, per_dataset={spectro.name: {"engine": "dag-ml"}})
def _separation_preproc_body_lists(steps: Any) -> list[list[Any]]:
"""Normalize shared-list or per-branch-dict preprocessing bodies."""
if isinstance(steps, list):
return [steps]
if isinstance(steps, dict) and steps and all(isinstance(value, list) for value in steps.values()):
return list(steps.values())
raise DagMlUnsupported("separation preprocessing concat requires branch steps as a list or per-branch dict of lists")
def _clean_separation_preproc_body(body: list[Any]) -> list[Any]:
"""Validate one branch preprocessing body and return X-transform steps."""
if any(isinstance(step, dict) for step in body):
raise DagMlUnsupported("separation preprocessing concat branch bodies may contain only X transforms")
cleaned = _supported_body_steps([step for step in body if not _is_split_step(step)])
for step in cleaned:
get_params = getattr(step, "get_params", None)
if not callable(get_params):
continue
try:
params = get_params(deep=False)
except Exception:
continue
if "random_state" in params and params["random_state"] is None:
raise DagMlUnsupported(
f"separation preprocessing concat requires {type(step).__name__}.random_state to be pinned"
)
return cleaned
def _separation_preproc_steps_for_branch(steps: Any, branch_name: str) -> list[Any]:
"""Resolve and validate preprocessing steps for one separation branch."""
from nirs4all.controllers.data.branch import BranchController
branch_steps = BranchController._resolve_branch_steps(steps, branch_name)
if not isinstance(branch_steps, list):
raise DagMlUnsupported(
f"separation preprocessing branch {branch_name!r} did not resolve to a list of steps"
)
return _clean_separation_preproc_body(branch_steps)
def _branch_filter_from_spec(filter_spec: Any) -> Any:
"""Deserialize a ``by_filter`` criterion into a SampleFilter instance."""
from nirs4all.operators.filters.native import resolve_sample_filter
filter_obj = resolve_sample_filter(filter_spec)
if filter_obj is None:
raise DagMlUnsupported(
f"by_filter preprocessing concat expects a SampleFilter, got {type(filter_spec).__name__}"
)
return _clone_operator_instance(filter_obj)
def _filter_y_arg(values: np.ndarray) -> np.ndarray:
"""SampleFilter y argument in the same flattened form used by tag/exclude helpers."""
return values.flatten() if values.ndim > 1 else values
def _tag_branch_values(pipeline: list[Any], spectro: Any, key: str, fit_pool: list[int]) -> dict[int, str]:
"""Compute the boolean tag values produced by the matching top-level tag step."""
from nirs4all.pipeline.dagml.steps import _taggers_from_step
taggers = [
(tag_name, filter_obj)
for step in pipeline
for tag_name, filter_obj in (_taggers_from_step(step) or [])
if tag_name == key
]
if len(taggers) != 1:
raise DagMlUnsupported(
f"by_tag preprocessing concat requires exactly one top-level tag step named {key!r}"
)
sample_ints = [int(sample_int) for sample_int in spectro.index_column("sample", {})]
if not sample_ints:
return {}
_tag_name, filter_obj = taggers[0]
filter_obj = _clone_operator_instance(filter_obj)
x_rows = np.asarray(spectro.x_rows(sample_ints, layout="2d"))
y_rows = _filter_y_arg(_dataset_y_rows(spectro, sample_ints))
filter_obj.fit(np.asarray(spectro.x_rows(fit_pool, layout="2d")), _filter_y_arg(_dataset_y_rows(spectro, fit_pool)))
mask = np.asarray(filter_obj.get_mask(x_rows, y_rows), dtype=bool)
if mask.shape[0] != len(sample_ints):
raise ValueError(f"tag filter {filter_obj.__class__.__name__} returned {mask.shape[0]} masks for {len(sample_ints)} samples")
return {sample_int: str(bool(not keep)) for sample_int, keep in zip(sample_ints, mask, strict=True)}
def _filter_branch_values(criterion: dict[str, Any], spectro: Any, train_pool: list[int]) -> dict[int, str]:
"""Compute ``passing``/``failing`` values for a ``by_filter`` branch criterion."""
names = criterion.get("names", ["passing", "failing"])
if not isinstance(names, list | tuple) or len(names) != 2:
raise DagMlUnsupported("by_filter preprocessing concat requires exactly two branch names")
filter_obj = _branch_filter_from_spec(criterion["by_filter"])
x_train = np.asarray(spectro.x_rows(train_pool, layout="2d"))
y_train = _filter_y_arg(_dataset_y_rows(spectro, train_pool))
filter_obj.fit(x_train, y_train)
sample_ints = [int(sample_int) for sample_int in spectro.index_column("sample", {})]
if not sample_ints:
return {}
x_rows = np.asarray(spectro.x_rows(sample_ints, layout="2d"))
y_rows = _filter_y_arg(_dataset_y_rows(spectro, sample_ints))
mask = np.asarray(filter_obj.get_mask(x_rows, y_rows), dtype=bool)
if mask.shape[0] != len(sample_ints):
raise ValueError(f"branch filter {filter_obj.__class__.__name__} returned {mask.shape[0]} masks for {len(sample_ints)} samples")
passing, failing = str(names[0]), str(names[1])
return {sample_int: passing if keep else failing for sample_int, keep in zip(sample_ints, mask, strict=True)}
def _separation_branch_values(
pipeline: list[Any],
spectro: Any,
criterion: dict[str, Any],
train_pool: list[int],
) -> dict[int, str]:
"""Branch-name value for every sample under a supported separation criterion."""
sample_ints = [int(sample_int) for sample_int in spectro.index_column("sample", {})]
if "by_metadata" in criterion:
key = str(criterion["by_metadata"])
values = spectro.metadata_column(key, {})
return {sample_int: str(value) for sample_int, value in zip(sample_ints, values, strict=True)}
if "by_tag" in criterion:
return _tag_branch_values(pipeline, spectro, str(criterion["by_tag"]), train_pool)
if "by_filter" in criterion:
return _filter_branch_values(criterion, spectro, train_pool)
raise DagMlUnsupported("separation preprocessing concat requires by_metadata, by_tag, or by_filter")
def _assign_separation_rows(
out: np.ndarray | None,
assigned_rows: set[int],
sample_pool: list[int],
row_samples: list[int],
current_rows: np.ndarray,
) -> np.ndarray:
"""Assign branch-projected rows into a pool-shaped matrix."""
if not sample_pool:
if out is not None:
return out
return np.empty((0, current_rows.shape[1]), dtype=current_rows.dtype)
if out is None:
out = np.empty((len(sample_pool), current_rows.shape[1]), dtype=current_rows.dtype)
elif out.shape[1] != current_rows.shape[1]:
raise DagMlUnsupported("separation preprocessing branches must produce the same feature width")
row_of = {sample_int: row for row, sample_int in enumerate(sample_pool)}
for position, sample_int in enumerate(row_samples):
out[row_of[int(sample_int)], :] = current_rows[position]
assigned_rows.add(row_of[int(sample_int)])
return out
def _assert_separation_rows_complete(out: np.ndarray | None, assigned_rows: set[int], sample_pool: list[int], width: int) -> np.ndarray:
"""Return a complete pool-shaped branch matrix, including the empty-pool case."""
if not sample_pool:
return np.empty((0, width), dtype=float)
if out is None:
raise DagMlUnsupported("separation preprocessing branch produced no feature rows")
missing_rows = set(range(len(sample_pool))) - assigned_rows
if missing_rows:
raise DagMlUnsupported(
"separation preprocessing concat cannot route every requested sample; "
f"{len(missing_rows)} row(s) had no train-backed branch"
)
return out
def _separation_preproc_train_test_matrices(
pipeline: list[Any],
spectro: Any,
*,
criterion: dict[str, Any],
steps: Any,
train_pool: list[int],
test_pool: list[int],
fit_pool: list[int] | None = None,
) -> tuple[np.ndarray, np.ndarray]:
"""Fit on the designated fold only, then reassemble all requested rows."""
fit_pool = train_pool if fit_pool is None else fit_pool
value_by_sample = _separation_branch_values(pipeline, spectro, criterion, fit_pool)
out_train: np.ndarray | None = None
out_test: np.ndarray | None = None
assigned_train: set[int] = set()
assigned_test: set[int] = set()
width = int(np.asarray(spectro.x_rows(train_pool, layout="2d")).shape[1])
branch_names = list(dict.fromkeys(value_by_sample[int(sample_int)] for sample_int in train_pool))
for branch_name in branch_names:
train_samples = [sample_int for sample_int in train_pool if value_by_sample[int(sample_int)] == branch_name]
test_samples = [sample_int for sample_int in test_pool if value_by_sample[int(sample_int)] == branch_name]
if not train_samples and not test_samples:
continue
body_steps = _separation_preproc_steps_for_branch(steps, branch_name)
# Membership selects rows at concat time, as in the public branch
# contract, but fitted state must observe only the current fold-train.
current_fit = np.asarray(spectro.x_rows(fit_pool, layout="2d"))
current_train = np.asarray(spectro.x_rows(train_samples, layout="2d"))
current_test = np.asarray(spectro.x_rows(test_samples, layout="2d")) if test_samples else np.empty((0, current_fit.shape[1]))
y_fit = _dataset_y_rows(spectro, fit_pool)
for step in body_steps:
transform = _clone_operator_instance(step)
transform.fit(current_fit, y_fit)
current_fit = np.asarray(transform.transform(current_fit))
current_train = np.asarray(transform.transform(current_train))
current_test = np.asarray(transform.transform(current_test)) if test_samples else np.empty((0, current_fit.shape[1]))
width = int(current_fit.shape[1])
if train_samples:
out_train = _assign_separation_rows(out_train, assigned_train, train_pool, train_samples, current_train)
if test_samples:
out_test = _assign_separation_rows(out_test, assigned_test, test_pool, test_samples, current_test)
return (
_assert_separation_rows_complete(out_train, assigned_train, train_pool, width),
_assert_separation_rows_complete(out_test, assigned_test, test_pool, width),
)
def _run_separation_preproc_concat(
pipeline: list[Any],
branch_step: dict[str, Any],
preproc_body: Any,
downstream_body: list[Any],
spectro: Any,
metric: str,
task_type: str,
config_name: str = "",
) -> RunResult:
"""Project separation preprocessing + concat + downstream model.
This path covers the legacy feature-reassembly shape where the branch applies
row-wise preprocessing per metadata/tag/filter partition, ``merge='concat'``
restores the full sample universe, and a single downstream model owns CV +
final refit rows.
"""
_, splitter = _split_pipeline(pipeline)
if splitter is None:
raise DagMlUnsupported("engine='dag-ml' requires a cross-validator step (e.g. KFold) in the pipeline")
if len(downstream_body) != 1:
raise DagMlUnsupported("separation preprocessing concat supports exactly one downstream model")
for body in _separation_preproc_body_lists(preproc_body):
_clean_separation_preproc_body(body)
downstream_steps = _apply_model_params(downstream_body)
_reject_multi_model(downstream_steps)
_assert_supported_operators([step for body in _separation_preproc_body_lists(preproc_body) for step in body] + downstream_steps)
model_step = downstream_steps[0]
if not isinstance(model_step, dict) or "model" not in model_step:
raise DagMlUnsupported("separation preprocessing concat requires a downstream model")
criterion = branch_step["branch"]
train_pool = [int(sample_int) for sample_int in spectro.index_column("sample", {"partition": "train"})]
test_pool = [int(sample_int) for sample_int in spectro.index_column("sample", {"partition": "test"})]
folds = _build_folds(splitter, spectro, train_pool, set())
X_train = np.asarray(spectro.x_rows(train_pool, layout="2d"))
X_test = np.asarray(spectro.x_rows(test_pool, layout="2d")) if test_pool else np.empty((0, X_train.shape[1]))
def matrix_for_fit(fit_pool: list[int]) -> tuple[np.ndarray, np.ndarray]:
return _separation_preproc_train_test_matrices(
pipeline, spectro, criterion=criterion, steps=preproc_body,
train_pool=train_pool, test_pool=test_pool, fit_pool=fit_pool,
)
return _run_model_on_precomputed_matrix(
X_train,
X_test,
_dataset_y_rows(spectro, train_pool),
_dataset_y_rows(spectro, test_pool) if test_pool else np.empty((0, 1)),
train_pool,
test_pool,
folds,
model_step["model"],
spectro,
metric,
task_type,
config_name,
branch_id=None,
branch_name="",
include_refit=True,
matrix_for_fit=matrix_for_fit,
)
def _as_prediction_matrix(values: Any) -> np.ndarray:
"""Model predictions as a 2D feature block."""
arr = np.asarray(values, dtype=float)
return arr.reshape(arr.shape[0], -1) if arr.ndim > 1 else arr.reshape(-1, 1)
def _branch_transform_operator(step: dict[str, Any]) -> Any:
"""Instantiate the structured branch transform dictionaries this compatibility path supports."""
if "concat_transform" not in step:
raise DagMlUnsupported(
"named MetaModel prediction-feature stack supports branch dict steps only for "
f"`concat_transform` or `model`, got {sorted(step)}"
)
from nirs4all.pipeline.dagml.operator_routing import route_operator
from nirs4all.pipeline.dagml_bridge import _lower_concat_transform
lowered = _lower_concat_transform(step)
return route_operator("transform", lowered["class"], lowered.get("params") or {})
def _with_fold_preprocessing(model: Any, preprocessing: list[Any] | None) -> Any:
"""Keep each fitted transform attached to its fold's estimator for all replay."""
if not preprocessing:
return model
from sklearn.pipeline import Pipeline
return Pipeline([*((f"transform_{index}", transform) for index, transform in enumerate(preprocessing)), ("model", model)])
def _prepare_named_branch_feature_matrix(branch: list[Any], spectro: Any, train_pool: list[int], test_pool: list[int]) -> tuple[dict[str, Any], np.ndarray, np.ndarray, list[Any]]:
"""Keep raw rows and unfitted branch transforms until the fold is selected."""
current_train = np.asarray(spectro.x_rows(train_pool, layout="2d"))
current_test = np.asarray(spectro.x_rows(test_pool, layout="2d")) if test_pool else np.empty((0, current_train.shape[1]))
model_steps: list[dict[str, Any]] = []
preprocessing: list[Any] = []
for step in _supported_body_steps([part for part in branch if not _is_split_step(part)]):
if isinstance(step, dict) and "model" in step:
model_steps.append(step)
continue
if model_steps:
raise DagMlUnsupported("named MetaModel prediction-feature stack does not support branch transforms after a branch model")
transform = _branch_transform_operator(step) if isinstance(step, dict) else _clone_operator_instance(step)
preprocessing.append(transform)
if len(model_steps) != 1:
raise DagMlUnsupported("named MetaModel prediction-feature stack requires exactly one model per branch")
return model_steps[0], current_train, current_test, preprocessing
def _project_cv_model_rows(
predictions: Predictions,
*,
X_train_all: np.ndarray,
X_test: np.ndarray,
train_pool: list[int],
test_pool: list[int],
folds: list[tuple[list[int], list[int]]],
model_template: Any,
model_name: str,
model_classname: str,
step_idx: int,
branch_id: int | None,
branch_name: str,
spectro: Any,
metric: str,
task_type: str,
config_name: str,
preprocessing: list[Any] | None = None,
) -> dict[str, Any]:
"""Emit legacy CV/avg/w_avg rows and return OOF/test features for downstream stacking."""
train_position = {sample_int: position for position, sample_int in enumerate(train_pool)}
model_template = _with_fold_preprocessing(model_template, preprocessing)
y_train_all = _dataset_y_rows(spectro, train_pool)
y_test = _dataset_y_rows(spectro, test_pool) if test_pool else np.empty((0, 1))
fold_models: list[Any] = []
fold_val_samples: list[list[int]] = []
fold_val_scores: list[float] = []
oof_predictions: np.ndarray | None = None
test_predictions: list[np.ndarray] = []
n_features = int(X_train_all.shape[1]) if X_train_all.ndim > 1 else 1
for fold_index, (fold_train, fold_val) in enumerate(folds):
train_rows = [train_position[int(sample_int)] for sample_int in fold_train]
val_rows = [train_position[int(sample_int)] for sample_int in fold_val]
model = _clone_operator_instance(model_template)
model.fit(X_train_all[train_rows], _dataset_y_rows(spectro, fold_train))
fold_models.append(model)
fold_val_samples.append([int(sample_int) for sample_int in fold_val])
pred_by_partition = {
"train": np.asarray(model.predict(X_train_all[train_rows])),
"val": np.asarray(model.predict(X_train_all[val_rows])),
"test": np.asarray(model.predict(X_test)) if len(test_pool) else np.array([]),
}
y_by_partition = {
"train": _dataset_y_rows(spectro, fold_train),
"val": _dataset_y_rows(spectro, fold_val),
"test": y_test,
}
scores = {partition: _score_block(y_by_partition[partition], pred_by_partition[partition], task_type) for partition in pred_by_partition if len(y_by_partition[partition]) and len(pred_by_partition[partition])}
fold_val_scores.append(float(scores["val"][metric]))
val_matrix = _as_prediction_matrix(pred_by_partition["val"])
if oof_predictions is None:
oof_predictions = np.full((len(train_pool), val_matrix.shape[1]), np.nan, dtype=float)
oof_predictions[val_rows, :] = val_matrix
if len(test_pool):
test_predictions.append(_as_prediction_matrix(pred_by_partition["test"]))
_add_scored_prediction_rows(
predictions,
spectro=spectro,
config_name=config_name,
step_idx=step_idx,
model_name=model_name,
model_classname=model_classname,
fold_id=str(fold_index),
scores=scores,
y_by_partition=y_by_partition,
pred_by_partition=pred_by_partition,
samples_by_partition={"train": [int(sample_int) for sample_int in fold_train], "val": [int(sample_int) for sample_int in fold_val], "test": test_pool},
metric=metric,
task_type=task_type,
n_features=n_features,
branch_id=branch_id,
branch_name=branch_name,
)
if oof_predictions is None:
raise DagMlUnsupported("named MetaModel prediction-feature stack produced no OOF predictions")
weights = EnsembleUtils._scores_to_weights(np.asarray(fold_val_scores, dtype=float), higher_is_better=is_higher_better(metric))
all_train_preds = [_as_prediction_matrix(model.predict(X_train_all)) for model in fold_models]
all_test_preds = [_as_prediction_matrix(model.predict(X_test)) for model in fold_models] if len(test_pool) else []
val_samples = [sample_int for fold_val in fold_val_samples for sample_int in fold_val]
val_rows_in_fold_order = [train_position[int(sample_int)] for sample_int in val_samples]
val_pred_oof = oof_predictions[val_rows_in_fold_order]
y_val_oof = np.vstack([_dataset_y_rows(spectro, fold_val) for fold_val in fold_val_samples])
selection_score = float(_score_block(y_val_oof, val_pred_oof, task_type)[metric])
avg_pred_by_partition = {
"train": np.mean(all_train_preds, axis=0),
"val": val_pred_oof,
"test": np.mean(all_test_preds, axis=0) if all_test_preds else np.array([]),
}
w_avg_pred_by_partition = {
"train": np.sum([weight * pred for weight, pred in zip(weights, all_train_preds, strict=False)], axis=0),
"val": val_pred_oof,
"test": np.sum([weight * pred for weight, pred in zip(weights, all_test_preds, strict=False)], axis=0) if all_test_preds else np.array([]),
}
y_by_partition = {"train": y_train_all, "val": y_val_oof, "test": y_test}
samples_by_partition = {"train": train_pool, "val": val_samples, "test": test_pool}
for fold_id, pred_by_partition, row_weights in (("avg", avg_pred_by_partition, None), ("w_avg", w_avg_pred_by_partition, weights)):
scores = {
partition: _score_block(y_by_partition[partition], pred_by_partition[partition], task_type)
for partition in pred_by_partition
if len(y_by_partition[partition]) and len(pred_by_partition[partition])
}
_add_scored_prediction_rows(
predictions,
spectro=spectro,
config_name=config_name,
step_idx=step_idx,
model_name=model_name,
model_classname=model_classname,
fold_id=fold_id,
scores=scores,
y_by_partition=y_by_partition,
pred_by_partition=pred_by_partition,
samples_by_partition=samples_by_partition,
metric=metric,
task_type=task_type,
n_features=n_features,
branch_id=branch_id,
branch_name=branch_name,
weights=row_weights,
)
test_mean = np.mean(all_test_preds, axis=0) if all_test_preds else np.empty((0, oof_predictions.shape[1]))
return {
"model_name": model_name,
"branch_id": branch_id,
"branch_name": branch_name,
"selection_score": selection_score,
"oof": oof_predictions,
"test": test_mean,
}
def _run_named_metamodel_feature_stack(
pipeline: list[Any],
branch_names: list[str],
branches: list[list[Any]],
meta_step: dict[str, Any],
prediction_configs: list[dict[str, Any]],
downstream_step: dict[str, Any],
spectro: Any,
metric: str,
task_type: str,
config_name: str = "",
) -> RunResult:
"""Project the legacy CV-only named-branch MetaModel → prediction-feature stack surface.
This path is intentionally a host-side compatibility projection: legacy fits named branch
preprocessing on the full train pool, trains branch-local models by CV, trains a branch-local
``MetaModel`` from each branch model's OOF predictions, selects the best model per requested branch
by validation RMSE, and finally trains a downstream estimator on those selected OOF features. The
legacy refit pass skips named-dict stacking, so the result contains no ``fold_id="final"`` rows.
"""
if task_type != "regression":
raise DagMlUnsupported("named MetaModel prediction-feature stack currently supports regression only")
_, splitter = _split_pipeline(pipeline)
if splitter is None:
raise DagMlUnsupported("engine='dag-ml' requires a cross-validator step (e.g. KFold) in the pipeline")
train_pool = [int(sample_int) for sample_int in spectro.index_column("sample", {"partition": "train"})]
test_pool = [int(sample_int) for sample_int in spectro.index_column("sample", {"partition": "test"})]
folds = _build_folds(splitter, spectro, train_pool, set())
branch_step_idx = next(index for index, step in enumerate(pipeline) if isinstance(step, dict) and "branch" in step) + 1
meta_step_idx = pipeline.index(meta_step) + 1
downstream_step_idx = pipeline.index(downstream_step) + 1
predictions = Predictions()
results_by_branch: dict[int, list[dict[str, Any]]] = {}
for branch_index, (branch_name, branch) in enumerate(zip(branch_names, branches, strict=True)):
model_step, X_branch_train, X_branch_test, preprocessing = _prepare_named_branch_feature_matrix(branch, spectro, train_pool, test_pool)
base_model_name = _model_name([model_step])
base_result = _project_cv_model_rows(
predictions,
X_train_all=X_branch_train,
X_test=X_branch_test,
train_pool=train_pool,
test_pool=test_pool,
folds=folds,
model_template=model_step["model"],
model_name=base_model_name,
preprocessing=preprocessing,
model_classname=type(model_step["model"]).__name__,
step_idx=branch_step_idx,
branch_id=branch_index,
branch_name=branch_name,
spectro=spectro,
metric=metric,
task_type=task_type,
config_name=config_name,
)
meta_operator = meta_step["model"]
meta_model_name = str(meta_step.get("name") or getattr(meta_operator, "name", None) or "MetaModel")
meta_result = _project_cv_model_rows(
predictions,
X_train_all=base_result["oof"],
X_test=base_result["test"],
train_pool=train_pool,
test_pool=test_pool,
folds=folds,
model_template=meta_operator.model,
model_name=meta_model_name,
model_classname=type(meta_operator).__name__,
step_idx=meta_step_idx,
branch_id=branch_index,
branch_name=branch_name,
spectro=spectro,
metric=metric,
task_type=task_type,
config_name=config_name,
)
results_by_branch[branch_index] = [base_result, meta_result]
selected_results: list[dict[str, Any]] = []
for prediction_config in prediction_configs:
branch_index = int(prediction_config["branch"])
candidates = sorted(results_by_branch[branch_index], key=lambda result: str(result["model_name"]))
reverse = is_higher_better(str(prediction_config.get("metric", metric)))
selected_results.append(sorted(candidates, key=lambda result: float(result["selection_score"]), reverse=reverse)[0])
X_downstream_train = np.concatenate([result["oof"] for result in selected_results], axis=1)
X_downstream_test = np.concatenate([result["test"] for result in selected_results], axis=1) if test_pool else np.empty((0, X_downstream_train.shape[1]))
_project_cv_model_rows(
predictions,
X_train_all=X_downstream_train,
X_test=X_downstream_test,
train_pool=train_pool,
test_pool=test_pool,
folds=folds,
model_template=downstream_step["model"],
model_name=_model_name([downstream_step]),
model_classname=type(downstream_step["model"]).__name__,
step_idx=downstream_step_idx,
branch_id=None,
branch_name="",
spectro=spectro,
metric=metric,
task_type=task_type,
config_name=config_name,
)
predictions.flush()
run_result = RunResult(predictions=predictions, per_dataset={spectro.name: {"engine": "dag-ml"}})
run_result._dagml_score_set = None # noqa: SLF001 - this path mirrors legacy rows without a Rust ScoreSet.
run_result._dagml_refit_artifacts = [] # noqa: SLF001
return run_result
def _run_duplication_merge_all_downstream_result(
branches: list[list[Any]],
splitter: Any,
model_step: dict[str, Any],
spectro: Any,
metric: str,
task_type: str,
config_name: str,
) -> RunResult:
"""Run downstream model rows over native-built ``merge='all'`` features + branch predictions."""
train_pool = [int(sample_int) for sample_int in spectro.index_column("sample", {"partition": "train"})]
test_pool = [int(sample_int) for sample_int in spectro.index_column("sample", {"partition": "test"})]
folds = _build_folds(splitter, spectro, train_pool, set())
train_position = {sample_int: position for position, sample_int in enumerate(train_pool)}
feature_train_parts: list[np.ndarray] = []
feature_test_parts: list[np.ndarray] = []
prediction_train_parts: list[np.ndarray] = []
prediction_test_parts: list[np.ndarray] = []
for branch in branches:
transforms, model_template, X_branch_train = _fit_global_branch_transforms(branch, spectro, train_pool)
X_branch_test = _transform_branch_matrix(spectro, test_pool, transforms) if test_pool else np.empty((0, X_branch_train.shape[1]))
feature_train_parts.append(X_branch_train)
feature_test_parts.append(X_branch_test)
oof_predictions = np.full((len(train_pool), 1), np.nan, dtype=float)
test_predictions: list[np.ndarray] = []
for fold_train, fold_val in folds:
train_rows = [train_position[int(sample_int)] for sample_int in fold_train]
val_rows = [train_position[int(sample_int)] for sample_int in fold_val]
model = _clone_operator_instance(model_template)
model.fit(X_branch_train[train_rows], _dataset_y_rows(spectro, fold_train))
oof_predictions[val_rows, :] = np.asarray(model.predict(X_branch_train[val_rows])).reshape(len(val_rows), -1)
if len(test_pool):
test_predictions.append(np.asarray(model.predict(X_branch_test)).reshape(len(test_pool), -1))
column_means = np.nanmean(oof_predictions, axis=0)
column_means = np.where(np.isnan(column_means), 0.0, column_means)
missing_rows, missing_cols = np.where(np.isnan(oof_predictions))
if len(missing_rows):
oof_predictions[missing_rows, missing_cols] = column_means[missing_cols]
prediction_train_parts.append(oof_predictions)
prediction_test_parts.append(np.mean(test_predictions, axis=0) if test_predictions else np.empty((0, oof_predictions.shape[1])))
X_train_all = np.concatenate([*feature_train_parts, *prediction_train_parts], axis=1)
X_test = np.concatenate([*feature_test_parts, *prediction_test_parts], axis=1) if test_pool else np.empty((0, X_train_all.shape[1]))
return _run_model_on_precomputed_matrix(
X_train_all,
X_test,
_dataset_y_rows(spectro, train_pool),
_dataset_y_rows(spectro, test_pool) if test_pool else np.empty((0, 1)),
train_pool,
test_pool,
folds,
model_step["model"],
spectro,
metric,
task_type,
config_name,
branch_id=None,
branch_name="",
)
def _run_by_source_concat_shared_preproc(pipeline: list[Any], preproc_body: list[Any], downstream_body: list[Any], n_sources: int, spectro: Any, dataset_arg: str, cli: str, venv_python: str, run_dir: Path, metric: str, task_type: str, dataset_pickle: str | None = None, config_name: str = "", random_state: int | None = None) -> RunResult:
"""Run shared by_source preprocessing + concat merge + one model as native CV/refit.
The downstream model is a normal single producer, but its compiled graph node carries
``metadata.source_concat_x_chain``. The process adapter then applies the upstream X-chain once per
feature source and hstacks those transformed blocks before fitting/predicting, matching legacy's
by_source feature reassembly for the shared-body LIST form. The final projection is replicated per
source because legacy stores the same downstream rows under each source branch.
"""
import dag_ml
if getattr(spectro, "independent_unit_ids", None) is not None:
raise DagMlUnsupported("experimental-unit influence does not support source-concat preprocessing")
_, splitter = _split_pipeline(pipeline)
if splitter is None:
raise DagMlUnsupported("engine='dag-ml' requires a cross-validator step (e.g. KFold) in the pipeline")
steps = _supported_body_steps([step for step in [*preproc_body, *downstream_body] if not _is_split_step(step)])
_reject_multi_model(steps)
_assert_supported_operators(steps)
steps = _apply_model_params(steps)
identity = mint_identity(spectro)
pool = spectro.index_column("sample", {"partition": "train"})
folds = _build_folds(splitter, spectro, pool, set())
envelope = build_envelope(spectro, identity, sample_ints=pool, group_by_sample=_split_group_grain(splitter, spectro, pool))
dsl = assemble_cv_refit_dsl(steps, identity, envelope, folds, dsl_id="nirs4all-by-source-concat", n_splits=len(folds))
dsl["root_seed"] = random_state if random_state is not None else 0
graph = dag_ml.compile_pipeline_dsl_artifact_with_controllers(dsl, controller_manifests()).graph.to_dict()
model_nodes = [node for node in graph["nodes"] if node["kind"] == "model"]
if len(model_nodes) != 1:
raise DagMlUnsupported(f"by_source concat compile produced {len(model_nodes)} model nodes, expected 1")
model_nodes[0]["metadata"] = {**model_nodes[0].get("metadata", {}), "source_concat_x_chain": True}
outcome = run_cv_refit_bundle(
dsl=dsl,
envelope=envelope,
graph=graph,
dataset_path=dataset_arg,
workdir=run_dir,
dagml_cli=cli,
venv_python=venv_python,
selection_metric=metric,
dataset_pickle=dataset_pickle,
dataset=spectro,
random_state=random_state,
)
if outcome["returncode"] != 0:
_raise_run_failure(outcome, "dag-ml by_source concat run failed")
result = _scores_to_run_result(
outcome["scores"],
spectro.name,
_model_name(steps),
metric,
task_type,
config_name=config_name,
results=outcome["results"],
identity=identity,
refit_artifacts=outcome["refit_artifacts"],
)
return _repeat_by_source_merge_projection(result, n_sources)
def _run_by_source_distinct_preproc_concat(
pipeline: list[Any],
source_steps: dict[str, list[Any]],
downstream_body: list[Any],
n_sources: int,
spectro: Any,
dataset_arg: str,
cli: str,
venv_python: str,
run_dir: Path,
metric: str,
task_type: str,
dataset_pickle: str | None = None,
config_name: str = "",
random_state: int | None = None,
) -> RunResult:
"""Run by_source DICT preprocessing + concat + downstream model as one native model node.
Each source's transform chain is cloned and fit on that source's fold-train block, the
transformed blocks are hstacked in ``source_layout.source_order``, and the downstream
model is fit on that concatenated matrix. Validation/test use the fitted per-source
chains, so preprocessing is fold-local and never fit on early-fused concat.
"""
import dag_ml
from nirs4all.pipeline.dagml.cli_runner import data_bindings_for_nodes, split_invocation_for
_, splitter = _split_pipeline(pipeline)
if splitter is None:
raise DagMlUnsupported("engine='dag-ml' requires a cross-validator step (e.g. KFold) in the pipeline")
if len(downstream_body) != 1:
raise DagMlUnsupported("by_source distinct preprocessing concat supports exactly one downstream model")
downstream_steps = _apply_model_params(downstream_body)
_reject_multi_model(downstream_steps)
model_node = _canonical_branch_step(downstream_steps[0], "model:source_concat")
if model_node["kind"] != "model":
raise DagMlUnsupported("by_source distinct preprocessing concat requires a downstream model")
identity = mint_identity(spectro)
pool = spectro.index_column("sample", {"partition": "train"})
folds = _build_folds(splitter, spectro, pool, set())
envelope = build_envelope(spectro, identity, sample_ints=pool, group_by_sample=_split_group_grain(splitter, spectro, pool))
source_layout = (envelope.get("plan") or {}).get("source_layout")
source_preprocessing = _source_preprocessing_metadata(source_steps, source_layout)
if len(source_preprocessing["sources"]) != n_sources:
raise DagMlUnsupported(
f"by_source distinct preprocessing source layout has {len(source_preprocessing['sources'])} source(s), expected {n_sources}"
)
model_node["metadata"] = {**model_node.get("metadata", {}), "source_concat_preprocessing": source_preprocessing}
canonical_dsl: dict[str, Any] = {
"id": "nirs4all-by-source-distinct-preproc-concat",
"steps": [model_node],
}
_bind_experimental_unit_design(canonical_dsl, spectro, identity, pool)
canonical_dsl["root_seed"] = random_state if random_state is not None else 0
graph = dag_ml.compile_pipeline_dsl_artifact_with_controllers(canonical_dsl, controller_manifests(canonical_dsl)).graph.to_dict()
model_ids = [node["id"] for node in graph["nodes"] if node["kind"] == "model"]
if model_ids != [model_node["id"]]:
raise DagMlUnsupported(f"by_source distinct preprocessing compile produced model nodes {model_ids!r}")
canonical_dsl["data_bindings"] = data_bindings_for_nodes(model_ids, envelope)
canonical_dsl["split_invocation"] = split_invocation_for(identity, folds, n_splits=len(folds))
outcome = run_cv_refit_bundle(
dsl=canonical_dsl,
envelope=envelope,
graph=graph,
dataset_path=dataset_arg,
workdir=run_dir,
dagml_cli=cli,
venv_python=venv_python,
selection_metric=metric,
dataset_pickle=dataset_pickle,
dataset=spectro,
random_state=random_state,
)
if outcome["returncode"] != 0:
_raise_run_failure(outcome, "dag-ml by_source distinct-preprocessing concat run failed")
result = _scores_to_run_result(
outcome["scores"],
spectro.name,
_model_name(downstream_steps),
metric,
task_type,
config_name=config_name,
results=outcome["results"],
identity=identity,
refit_artifacts=outcome["refit_artifacts"],
)
return _repeat_by_source_merge_projection(result, n_sources)
def _run_by_source_branch(pipeline: list[Any], branch_body: list[Any], aggregate: str, n_sources: int, spectro: Any, dataset_arg: str, cli: str, venv_python: str, run_dir: Path, metric: str, task_type: str, dataset_pickle: str | None = None, config_name: str = "", random_state: int | None = None) -> RunResult:
"""Run a by_source separation branch + avg/mean fusion merge as ONE native dag-ml run (S4).
LATE fusion BY SOURCE: fans the shared body into one canonical branch PER feature source
(:func:`_canonical_source_branch`), each MODEL node bound to its source via ``metadata.source_index``
so the node runner feeds it ONLY that source's block (all samples, that source's columns). The fold
set, OOF, and merge are sample-keyed and identical to the duplication-fusion path — the ONLY
difference from duplication is the feature-axis (per-source) restriction each branch's model sees.
dag-ml runs ONE native CV+refit: each per-source branch model is FIT_CV on the full fold-train
(its source's columns) and predicts the full fold-validation (held-out OOF); the native fusion
merge handler averages the branches' per-sample Validation OOF (leakage-safe — train predictions
never enter the average) into one full-universe OOF attributed to the merge node, whose cross-fold
average is ``cv_best_score``. ``best_rmse`` (final test) is also native: each branch's REFIT predicts
the held-out TEST set (its source's columns, ``fold_id=None``) and dag-ml's off-fold fusion handler
averages those per sample into one scored ``(test, fold_id=None)`` block under the merge node.
LEAKAGE: folds/OOF over SAMPLES (unchanged — all branches see all samples, just different source
columns); a sample's blocks all land on the same fold side (the source restriction is a feature-axis
selection, never a sample partition); per-source per-fold preprocessing fits on fold-train only
(the per-branch X-chain is fit inside the fold's train materialization, like the single-block path);
the fusion merge is OOF-safe. No cross-sample mixing.
Classification (``fusion_proba_mean``) is NOT wired (identical to the duplication-fusion path): the
node runner emits scalar value predictions, not per-class probability rows, so a probability-mean
fusion has no proba blocks to average — it fails loud rather than silently running ``proba_mean`` as
a value (regression) fusion (audit H-P0-1).
"""
import dag_ml
from nirs4all.pipeline.dagml.cli_runner import data_bindings_for_nodes, split_invocation_for
if aggregate == "proba_mean":
raise DagMlUnsupported(
"engine='dag-ml' does not yet support proba-mean fusion (classification) for a by_source branch: "
"the process adapter emits class-label predictions, not per-class probability rows; backlog #20-avg (proba)."
)
_, splitter = _split_pipeline(pipeline)
if splitter is None:
raise DagMlUnsupported("engine='dag-ml' requires a cross-validator step (e.g. KFold) in the pipeline")
identity = mint_identity(spectro)
# The handled shape rejects any exclude step, so the CV universe is the full train pool.
pool = spectro.index_column("sample", {"partition": "train"})
folds = _build_folds(splitter, spectro, pool, set())
envelope = build_envelope(spectro, identity, sample_ints=pool, group_by_sample=_split_group_grain(splitter, spectro, pool))
# Canonical DSL: one duplication-mode branch with N per-source sub-pipelines (each the shared body
# bound to its source via metadata.source_index) + a fusion merge. The branch is `duplication` mode
# because every branch sees the FULL fold sample view (no fan-out / no branch_view / no
# sample_metadata) — the per-source restriction is applied host-side at materialization, not by a
# sample-partition branch_view.
canonical_dsl: dict[str, Any] = {
"id": "nirs4all-by-source-fusion",
"steps": [
{"kind": "branch", "mode": "duplication", "branches": [_canonical_source_branch(branch_body, source_index) for source_index in range(n_sources)]},
{"kind": "merge", "id": _FUSION_MERGE_NODE_ID, "merge_mode": "fusion", "output_as": "predictions"},
],
}
manifests = controller_manifests()
canonical_dsl["root_seed"] = random_state if random_state is not None else 0
graph = dag_ml.compile_pipeline_dsl_artifact_with_controllers(canonical_dsl, manifests).graph.to_dict()
model_ids = [node["id"] for node in graph["nodes"] if node["kind"] == "model"]
if len(model_ids) != n_sources:
raise DagMlUnsupported(f"by_source compile produced {len(model_ids)} model nodes, expected {n_sources}")
# One data_binding per per-source model node (each binds its `x` to the full source set; the node
# runner selects its own source's block by metadata.source_index) + the materialized fold set.
canonical_dsl["data_bindings"] = data_bindings_for_nodes(model_ids, envelope)
canonical_dsl["split_invocation"] = split_invocation_for(identity, folds, n_splits=len(folds))
outcome = run_cv_refit_bundle(
dsl=canonical_dsl, envelope=envelope, graph=graph, dataset_path=dataset_arg, workdir=run_dir, dagml_cli=cli, venv_python=venv_python, selection_metric=metric, dataset_pickle=dataset_pickle, dataset=spectro, random_state=random_state
)
if outcome["returncode"] != 0:
_raise_run_failure(outcome, "dag-ml by_source run failed")
model_label = f"by_source_{_model_name(branch_body)}x{n_sources}"
return _scores_to_run_result(outcome["scores"], spectro.name, model_label, metric, task_type, producer=_FUSION_MERGE_NODE_ID, config_name=config_name, refit_artifacts=outcome["refit_artifacts"])
def _source_names(spectro: Any, n_sources: int) -> list[str]:
names: list[str] = []
for source_index in range(n_sources):
name = None
if hasattr(spectro, "source_name"):
try:
name = spectro.source_name(source_index)
except Exception: # noqa: BLE001 - optional legacy dataset hook.
name = None
names.append(str(name) if name else f"source_{source_index}")
return names
def _run_rep_to_sources_by_source(
pipeline: list[Any], rep_step: dict[str, Any], branch_body: list[Any], spectro: Any,
dataset_arg: str, cli: str, venv_python: str, run_dir: Path, metric: str,
task_type: str, config_name: str = "", random_state: int | None = None,
) -> RunResult:
"""Materialize the legacy repetition reshape, then run native source-local models."""
import copy
import pickle
reshaped = copy.deepcopy(spectro)
_reshape_for_rep_fusion(rep_step, reshaped)
n_sources = reshaped.features_sources()
if n_sources < 2:
raise DagMlUnsupported("rep_to_sources by_source requires at least two aligned sources")
run_dir.mkdir(parents=True, exist_ok=True)
pickle_path = run_dir / "reshaped_dataset.pkl"
pickle_path.write_bytes(pickle.dumps(reshaped))
source_bodies = dict.fromkeys(_source_names(reshaped, n_sources), branch_body)
body_pipeline = [step for step in pipeline if step is not rep_step]
return _run_by_source_auto_models(
body_pipeline, source_bodies, [], n_sources, reshaped, dataset_arg, cli,
venv_python, run_dir / "source_models", metric, task_type,
dataset_pickle=str(pickle_path), config_name=config_name,
random_state=random_state,
)
def _run_by_source_auto_models(
pipeline: list[Any], source_bodies: dict[str, list[Any]], y_steps: list[Any],
n_sources: int, spectro: Any, dataset_arg: str, cli: str, venv_python: str,
run_dir: Path, metric: str, task_type: str, dataset_pickle: str | None = None,
config_name: str = "", random_state: int | None = None,
refit_top_k: int = 1,
refit: bool = True,
) -> RunResult:
"""Score every source-local model when auto merge has no downstream estimator."""
import dag_ml
from nirs4all.pipeline.dagml.cli_runner import data_bindings_for_nodes, split_invocation_for
names = _source_names(spectro, n_sources)
if set(source_bodies) != set(names):
raise DagMlUnsupported(f"by_source model names {list(source_bodies)!r} must match source names {names!r}")
_, splitter = _split_pipeline(pipeline)
if splitter is None:
raise DagMlUnsupported("by_source per-source model comparison requires a cross-validator")
identity = mint_identity(spectro)
pool = spectro.index_column("sample", {"partition": "train"})
folds = _build_folds(splitter, spectro, pool, set())
envelope = build_envelope(spectro, identity, sample_ints=pool, group_by_sample=_split_group_grain(splitter, spectro, pool))
has_operator_generator = any(
isinstance(step, dict) and "_or_" in step
for body in source_bodies.values() for step in body
)
branches = [
_canonical_source_branch([*y_steps, *source_bodies[name]], index, force_generator=has_operator_generator)
for index, name in enumerate(names)
]
canonical_dsl: dict[str, Any] = {
"id": "nirs4all-by-source-auto-models",
"steps": [{"kind": "branch", "mode": "duplication", "branches": branches}],
}
from .attested_by_source import bind_source_output_contract, execute_attested_by_source_cv
bind_source_output_contract(canonical_dsl, spectro, identity, pool, names)
canonical_dsl["root_seed"] = random_state if random_state is not None else 0
graph = dag_ml.compile_pipeline_dsl_artifact_with_controllers(canonical_dsl, controller_manifests(canonical_dsl)).graph.to_dict()
model_nodes = [node for node in graph["nodes"] if node["kind"] == "model"]
model_ids = [node["id"] for node in model_nodes]
compiled_model_ids = {node["id"] for node in graph["nodes"] if node["kind"] == "model"}
if len(compiled_model_ids) < n_sources or len(model_ids) != len(compiled_model_ids):
raise DagMlUnsupported(f"by_source auto model compile produced {compiled_model_ids!r}, expected {model_ids!r}")
canonical_dsl["data_bindings"] = data_bindings_for_nodes(model_ids, envelope)
canonical_dsl["split_invocation"] = split_invocation_for(identity, folds, n_splits=len(folds))
from .in_process_runner import _dagml_extension_loads, in_process_enabled
use_in_process = in_process_enabled() and _dagml_extension_loads()
attested = execute_attested_by_source_cv(
dsl=canonical_dsl, envelope=envelope, graph=graph, spectro=spectro,
identity=identity, folds=folds, source_names=names,
selection_metric=metric, random_state=random_state, refit_top_k=refit_top_k, refit=refit,
group_by_sample=_split_group_grain(splitter, spectro, pool),
cli=None if use_in_process else cli, python=venv_python, dataset_path=dataset_arg,
dataset_pickle=dataset_pickle, workdir=run_dir,
)
outcome = {"returncode": 0, "scores": attested["scores"], "results": attested["results"],
"refit_artifacts": attested["refit_artifacts"], "selected_refit_variant_ids": attested["selected_refit_variant_ids"]}
if outcome["returncode"] != 0:
_raise_run_failure(outcome, "dag-ml by_source auto model run failed")
winner_variant_id = next(
(
report.get("variant_id")
for report in (outcome["scores"] or {}).get("reports", [])
if report.get("partition") == "final" and report.get("fold_id") is None
),
None,
)
if winner_variant_id is None and not refit:
winner_variant_id = next(
(report.get("variant_id") for report in (outcome["scores"] or {}).get("reports", [])
if report.get("partition") == "validation" and report.get("fold_id") != "avg"),
None,
)
results_by_variant = _frames_by_variant(outcome["results"], winner_variant_id) if winner_variant_id is not None else None
predictions = Predictions()
for node in model_nodes:
model_id = node["id"]
index = (node.get("metadata") or {}).get("source_index")
if not isinstance(index, int) or not 0 <= index < n_sources:
raise DagMlUnsupported(f"by_source model {model_id!r} has no valid native source_index")
name = names[index]
local = _scores_to_run_result(
outcome["scores"], spectro.name, _model_name(source_bodies[name]), metric,
task_type, producer=model_id, config_name=config_name,
results_by_variant=results_by_variant, identity=identity,
refit_artifacts=outcome["refit_artifacts"],
emit_all_refits=refit_top_k > 1,
refit_name_suffix=f"_refit_rmsecvt{refit_top_k}" if refit_top_k > 1 else "_refit",
)
for row in local.predictions.filter_predictions(load_arrays=True):
row["branch_id"] = index
row["branch_name"] = name
predictions.extend_from_list([row])
predictions.flush()
result = RunResult(predictions=predictions, per_dataset={spectro.name: {
"engine": "dag-ml", "output_topology": "independent_by_source",
}})
result._dagml_score_set = outcome["scores"] # noqa: SLF001
result._dagml_node_results = outcome["results"] # noqa: SLF001
result._dagml_refit_artifacts = outcome["refit_artifacts"] # noqa: SLF001
from .envelope import _numeric_feature_axis
result._dagml_source_feature_axes = tuple(_numeric_feature_axis(spectro, index) for index in range(n_sources)) # noqa: SLF001
if attested is not None:
result._dagml_training_outcome = attested["training_outcome"] # noqa: SLF001
result._dagml_portable_predictor_package = attested["portable_package_dict"] # noqa: SLF001
result._dagml_source_training_captures = attested["captures"] # noqa: SLF001
if refit_top_k > 1:
result.per_dataset[spectro.name]["selected_refit_variant_ids"] = outcome.get("selected_refit_variant_ids", [])
return result
def _run_by_source_stacking_branch(
pipeline: list[Any],
branch_body: list[Any] | dict[str, list[Any]],
meta_learner: Any,
n_sources: int,
spectro: Any,
dataset_arg: str,
cli: str,
venv_python: str,
run_dir: Path,
metric: str,
task_type: str,
dataset_pickle: str | None = None,
config_name: str = "",
random_state: int | None = None,
) -> RunResult:
"""Execute actual per-source prediction stacking through native nested OOF.
Source slices precede each branch's preprocessing, fitted inside each
native training view. The meta-model consumes predictions, never duplicated
spectral blocks. Explicit feature/source concatenation is a separate path.
"""
from nirs4all.data.multimodal import MultimodalSpectroDataset
from .source_stacking import lower_source_stacking
widths = spectro.num_features
if not isinstance(widths, list) or len(widths) != n_sources:
raise DagMlUnsupported("source stacking requires aligned per-source feature counts")
lowered, branches, layout = lower_source_stacking(
pipeline, branch_body, source_widths=[int(width) for width in widths],
source_names=_source_names(spectro, n_sources),
source_descriptors=spectro.cohort.schema_descriptors() if isinstance(spectro, MultimodalSpectroDataset) else None,
)
return _run_stacking_branch(
lowered, branches, meta_learner, spectro, dataset_arg, cli, venv_python,
run_dir, metric, task_type, dataset_pickle=dataset_pickle,
config_name=config_name, random_state=random_state, source_layout=layout,
)
def _run_duplication_branch_feature_merge(pipeline: list[Any], branches: list[list[Any]], merge_mode: str, spectro: Any, dataset_arg: str, cli: str, venv_python: str, run_dir: Path, metric: str, task_type: str, dataset_pickle: str | None, config_name: str, random_state: int | None, refit: bool = True) -> RunResult:
"""Run legacy duplication ``merge=features``/``merge=all`` through one concrete native model."""
if merge_mode not in ("features", "all"):
raise DagMlUnsupported("engine='dag-ml' supports duplication branch feature merge only for merge='features' or merge='all'")
_, splitter = _split_pipeline(pipeline)
model_step = next((step for step in pipeline if isinstance(step, dict) and "model" in step), None)
if splitter is None or model_step is None:
raise DagMlUnsupported(f"engine='dag-ml' requires splitter + downstream model for duplication merge={merge_mode!r}")
if merge_mode == "features":
transformer = _branch_merge_transformer_step(branches, merge_mode)
synthetic_pipeline = [splitter, transformer, model_step]
downstream_result = _run_concrete(
synthetic_pipeline, spectro, dataset_arg, cli, venv_python, run_dir / "downstream", metric, task_type, dataset_pickle=dataset_pickle, config_name=config_name, random_state=random_state, refit=refit
)
return downstream_result
branch_names = _duplication_branch_names(pipeline, len(branches))
branch_results = [
_run_duplication_merge_all_branch_result(branch, index, branch_names[index], splitter, spectro, metric, task_type, config_name)
for index, branch in enumerate(branches)
]
downstream_result = _run_duplication_merge_all_downstream_result(branches, splitter, model_step, spectro, metric, task_type, config_name)
return _combine_duplication_merge_all_rows(branch_results, branch_names, downstream_result, _model_name([model_step]), spectro.name)
def _run_duplication_branch(pipeline: list[Any], branches: list[list[Any]], aggregate: str, spectro: Any, dataset_arg: str, cli: str, venv_python: str, run_dir: Path, metric: str, task_type: str, dataset_pickle: str | None = None, config_name: str = "", random_state: int | None = None, refit: bool = True) -> RunResult:
"""Run a duplication branch (``[[A], [B], …]``) + avg/mean fusion merge as ONE native dag-ml run.
Lowers each inner sub-pipeline to a canonical branch (``mode: "duplication"`` — every branch model
node gets the FULL fold data view: NO fan-out, NO ``auto_separate``, NO ``branch_view``,
NO ``sample_metadata``) and a fusion merge node (``merge_mode: "fusion"`` for the value mean,
``output_as: "predictions"``). dag-ml runs ONE native CV+refit: each branch model is FIT_CV on the
full fold-train and predicts the full fold-validation (held-out OOF); the native fusion merge handler
averages the branches' per-sample Validation OOF (leakage-safe — train predictions never enter the
average) into one full-universe OOF attributed to the merge node, whose cross-fold average is
``cv_best_score``.
``best_rmse`` (final test) is also native: each branch's REFIT predicts the held-out TEST set
(the node runner emits it with ``fold_id=None``), and dag-ml's off-fold fusion handler
(``reassemble_branch_merge_off_fold``) averages those base test predictions per sample into one
scored ``(test, fold_id=None)`` block under the merge node — the same average as ``cv_best_score``
but over the held-out test set.
Classification (``fusion_proba_mean``) is NOT wired: the node runner emits scalar value predictions,
not per-class probability rows, so a probability-mean fusion has no proba blocks to average — it
fails loud rather than averaging class labels (which is not what ``proba_mean`` means).
"""
if aggregate in ("features", "all"):
return _run_duplication_branch_feature_merge(
pipeline, branches, aggregate, spectro, dataset_arg, cli, venv_python, run_dir, metric, task_type, dataset_pickle, config_name, random_state, refit=refit
)
import dag_ml
from nirs4all.pipeline.dagml.cli_runner import data_bindings_for_nodes, split_invocation_for
if aggregate == "proba_mean":
raise NotImplementedError(
"engine='dag-ml' does not yet support proba-mean fusion (classification): the process adapter "
"emits class-label predictions, not per-class probability rows; backlog #20-avg (proba)."
)
merge_mode = "fusion"
_, splitter = _split_pipeline(pipeline)
if splitter is None:
raise DagMlUnsupported("engine='dag-ml' requires a cross-validator step (e.g. KFold) in the pipeline")
identity = mint_identity(spectro)
# The handled shape rejects any exclude step, so the CV universe is the full train pool.
pool = spectro.index_column("sample", {"partition": "train"})
folds = _build_folds(splitter, spectro, pool, set())
envelope = build_envelope(spectro, identity, sample_ints=pool, group_by_sample=_split_group_grain(splitter, spectro, pool))
# Canonical DSL: one duplication branch with N sub-pipelines (each on the FULL data) + a fusion merge.
canonical_dsl: dict[str, Any] = {
"id": "nirs4all-duplication-fusion",
"steps": [
{"kind": "branch", "mode": "duplication", "branches": [_canonical_branch(branch, index) for index, branch in enumerate(branches)]},
{"kind": "merge", "id": _FUSION_MERGE_NODE_ID, "merge_mode": merge_mode, "output_as": "predictions"},
],
}
manifests = controller_manifests()
canonical_dsl["root_seed"] = random_state if random_state is not None else 0
graph = dag_ml.compile_pipeline_dsl_artifact_with_controllers(canonical_dsl, manifests).graph.to_dict()
model_ids = [node["id"] for node in graph["nodes"] if node["kind"] == "model"]
if len(model_ids) < 2:
raise DagMlUnsupported("duplication-fusion compile produced fewer than two model nodes")
# One data_binding per branch model node (each binds its `x` to the full source) + the materialized
# fold set. Every model node sees the full fold view — no branch_view/sample_metadata (duplication).
canonical_dsl["data_bindings"] = data_bindings_for_nodes(model_ids, envelope)
canonical_dsl["split_invocation"] = split_invocation_for(identity, folds, n_splits=len(folds))
outcome = run_cv_refit_bundle(
dsl=canonical_dsl, envelope=envelope, graph=graph, dataset_path=dataset_arg, workdir=run_dir, dagml_cli=cli, venv_python=venv_python, selection_metric=metric, dataset_pickle=dataset_pickle, dataset=spectro, random_state=random_state
)
if outcome["returncode"] != 0:
_raise_run_failure(outcome, "dag-ml duplication-fusion run failed")
# The fusion-merge producer's reports carry the full-universe cross-fold OOF average (the fused
# ensemble's `cv_best_score`) AND a reassembled `(test, fold_id=None)` block (`best_rmse`, the
# branches' test predictions averaged per sample). Both are surfaced by `_scores_to_run_result`.
model_label = "+".join(_model_name(branch) for branch in branches)
return _scores_to_run_result(outcome["scores"], spectro.name, model_label, metric, task_type, producer=_FUSION_MERGE_NODE_ID, config_name=config_name, refit_artifacts=outcome["refit_artifacts"])
_META_NODE_ID = "merge:stack"
def _uses_named_duplication_branch(pipeline: list[Any]) -> bool:
"""Whether the original stacking branch used the legacy named-dict duplication syntax."""
branch_step = next((step for step in pipeline if isinstance(step, dict) and "branch" in step), None)
if not isinstance(branch_step, dict):
return False
raw_branch = branch_step.get("branch")
return isinstance(raw_branch, dict) and all(isinstance(body, list) for body in raw_branch.values())
def _score_prediction_block(y_true: Any, y_pred: Any, task_type: str) -> dict[str, float]:
"""Evaluate one prediction block using the same metric bundle as legacy prediction rows."""
from nirs4all.core import metrics as evaluator
values = evaluator.eval_multi(np.asarray(y_true, dtype=float).ravel(), np.asarray(y_pred, dtype=float).ravel(), task_type)
return {name: float(value) for name, value in values.items()}
def _branch_sklearn_pipeline(branch: list[Any]) -> Any:
"""Clone a branch body into a sklearn estimator pipeline for compatibility-row scoring."""
from sklearn.base import clone
from sklearn.pipeline import make_pipeline
operators: list[Any] = []
for step in _supported_body_steps([part for part in branch if not _is_split_step(part)]):
if isinstance(step, dict) and "model" in step:
operators.append(clone(step["model"]))
elif isinstance(step, dict):
raise DagMlUnsupported(f"named stacking projection does not support structured branch step {sorted(step)!r}")
else:
operators.append(clone(step))
if not operators:
raise DagMlUnsupported("named stacking projection requires non-empty branch bodies")
return make_pipeline(*operators) if len(operators) > 1 else operators[0]
def _mean_prediction_score(y_true: np.ndarray, predictions: list[np.ndarray], task_type: str) -> dict[str, float]:
"""Score the per-sample mean of a fold-model prediction ensemble."""
return _score_prediction_block(y_true, np.mean(np.stack(predictions, axis=0), axis=0), task_type)
def _emit_named_stacking_row(
predictions: Predictions,
*,
dataset_name: str,
config_name: str,
model_name: str,
fold_id: str,
partition: str,
metric: str,
task_type: str,
scores: dict[str, dict[str, float]],
branch_id: int | None = None,
) -> None:
"""Add one score-only row in the legacy branch-stacking table shape."""
predictions.add_prediction(
dataset_name=dataset_name,
config_name=config_name,
model_name=model_name,
fold_id=fold_id,
partition=partition,
metric=metric,
task_type=task_type,
scores=scores,
train_score=scores.get("train", {}).get(metric),
val_score=scores.get("val", {}).get(metric),
test_score=scores.get("test", {}).get(metric),
branch_id=branch_id,
)
def _cv_only_stacking_score_set(scores: dict[str, Any] | None) -> dict[str, Any] | None:
"""Return a ScoreSet with only CV validation reports for legacy no-refit stacking surfaces."""
if scores is None:
return None
projected = dict(scores)
projected["reports"] = [report for report in scores.get("reports", []) if report.get("partition") == "validation"]
return projected
def _named_dict_stacking_legacy_projection(
*,
pipeline: list[Any],
branches: list[list[Any]],
meta_learner: Any,
spectro: Any,
folds: list[tuple[list[int], list[int]]],
metric: str,
task_type: str,
config_name: str,
scores: dict[str, Any] | None,
) -> RunResult:
"""Project named-dict duplication stacking as legacy's CV-only branch table.
Legacy treats a named-dict duplication branch followed by ``{"merge": "predictions"}`` as stacking
during CV, but its refit pass only recognizes list-form duplication branches and therefore skips the
final refit. The public table is consequently 3 base branch models plus the downstream Ridge
meta-learner, each with 3 folds x train/val/test and avg/w_avg x train/val/test rows, and no
``fold_id="final"`` rows. This historical oracle is retained for its dedicated regression tests;
production dispatch never calls it. Named and list-form stacking now share native nested
scheduling and captured final artifacts instead of this historical local execution.
The meta rows are computed only from validation OOF branch predictions. Held-out test branch
predictions are used as prediction inputs for scoring, never as training data.
"""
from sklearn.base import clone
from sklearn.model_selection import LeaveOneOut, cross_val_predict
if task_type != "regression":
raise DagMlUnsupported("named stacking projection currently supports regression parity only")
pool = [int(sample) for sample in spectro.index_column("sample", {"partition": "train"})]
test_ids = [int(sample) for sample in spectro.index_column("sample", {"partition": "test"})]
def x(sample_ids: list[int]) -> np.ndarray:
return np.asarray(spectro.x({"sample": sample_ids}, layout="2d"))
def y(sample_ids: list[int]) -> np.ndarray:
return np.asarray(spectro.y({"sample": sample_ids}), dtype=float).ravel()
y_pool = y(pool)
y_test = y(test_ids)
branch_records: list[dict[str, Any]] = []
for branch_index, branch in enumerate(branches):
fold_records: list[dict[str, Any]] = []
pool_predictions: list[np.ndarray] = []
test_predictions: list[np.ndarray] = []
val_true: list[np.ndarray] = []
val_pred: list[np.ndarray] = []
for train_ids, val_ids in folds:
estimator = _branch_sklearn_pipeline(branch)
estimator.fit(x(train_ids), y(train_ids))
record: dict[str, Any] = {}
for partition, sample_ids in (("train", train_ids), ("val", val_ids), ("test", test_ids)):
y_true = y(sample_ids)
y_pred = np.asarray(estimator.predict(x(sample_ids)), dtype=float).ravel()
record[partition] = {"y_true": y_true, "y_pred": y_pred, "scores": _score_prediction_block(y_true, y_pred, task_type)}
fold_records.append(record)
pool_predictions.append(np.asarray(estimator.predict(x(pool)), dtype=float).ravel())
test_predictions.append(record["test"]["y_pred"])
val_true.append(record["val"]["y_true"])
val_pred.append(record["val"]["y_pred"])
avg_scores = {
"train": _mean_prediction_score(y_pool, pool_predictions, task_type),
"val": _score_prediction_block(np.concatenate(val_true), np.concatenate(val_pred), task_type),
"test": _mean_prediction_score(y_test, test_predictions, task_type),
}
branch_records.append(
{
"branch_id": branch_index,
"model_name": _model_name(branch),
"folds": fold_records,
"avg_scores": avg_scores,
"pool_predictions": pool_predictions,
"test_predictions": test_predictions,
}
)
meta_fold_records: list[dict[str, Any]] = []
meta_pool_predictions: list[np.ndarray] = []
meta_test_predictions: list[np.ndarray] = []
meta_val_true: list[np.ndarray] = []
meta_val_pred: list[np.ndarray] = []
for fold_index, (train_ids, val_ids) in enumerate(folds):
x_meta_val = np.column_stack([record["folds"][fold_index]["val"]["y_pred"] for record in branch_records])
x_meta_train = np.column_stack([record["folds"][fold_index]["train"]["y_pred"] for record in branch_records])
x_meta_test = np.column_stack([record["folds"][fold_index]["test"]["y_pred"] for record in branch_records])
x_meta_pool = np.column_stack([record["pool_predictions"][fold_index] for record in branch_records])
y_meta_val = y(val_ids)
direct_meta = clone(meta_learner)
direct_meta.fit(x_meta_val, y_meta_val)
# Validation rows must remain out-of-fold. Substituting an in-sample
# direct_meta.predict(x_meta_val) score here leaks the validation target
# into the public leaderboard and can promote the meta-learner unfairly.
if len(y_meta_val) > 1:
meta_fold_val_pred = np.asarray(cross_val_predict(clone(meta_learner), x_meta_val, y_meta_val, cv=LeaveOneOut()), dtype=float).ravel()
else:
meta_fold_val_pred = np.asarray(direct_meta.predict(x_meta_val), dtype=float).ravel()
train_pred = np.asarray(direct_meta.predict(x_meta_train), dtype=float).ravel()
test_pred = np.asarray(direct_meta.predict(x_meta_test), dtype=float).ravel()
pool_pred = np.asarray(direct_meta.predict(x_meta_pool), dtype=float).ravel()
record = {
"train": {"y_true": y(train_ids), "y_pred": train_pred, "scores": _score_prediction_block(y(train_ids), train_pred, task_type)},
"val": {"y_true": y_meta_val, "y_pred": meta_fold_val_pred, "scores": _score_prediction_block(y_meta_val, meta_fold_val_pred, task_type)},
"test": {"y_true": y_test, "y_pred": test_pred, "scores": _score_prediction_block(y_test, test_pred, task_type)},
}
meta_fold_records.append(record)
meta_pool_predictions.append(pool_pred)
meta_test_predictions.append(test_pred)
meta_val_true.append(y_meta_val)
meta_val_pred.append(meta_fold_val_pred)
meta_avg_scores = {
"train": _mean_prediction_score(y_pool, meta_pool_predictions, task_type),
"val": _score_prediction_block(np.concatenate(meta_val_true), np.concatenate(meta_val_pred), task_type),
"test": _mean_prediction_score(y_test, meta_test_predictions, task_type),
}
projected = Predictions()
all_records = [
*branch_records,
{
"branch_id": None,
"model_name": type(meta_learner).__name__,
"folds": meta_fold_records,
"avg_scores": meta_avg_scores,
},
]
for model_record in all_records:
for fold_index, fold_record in enumerate(model_record["folds"]):
fold_scores = {partition: fold_record[partition]["scores"] for partition in ("train", "val", "test")}
for partition in ("train", "val", "test"):
_emit_named_stacking_row(
projected,
dataset_name=spectro.name,
config_name=config_name,
model_name=model_record["model_name"],
fold_id=str(fold_index),
partition=partition,
metric=metric,
task_type=task_type,
scores=fold_scores,
branch_id=model_record["branch_id"],
)
for fold_id in ("avg", "w_avg"):
for partition in ("train", "val", "test"):
_emit_named_stacking_row(
projected,
dataset_name=spectro.name,
config_name=config_name,
model_name=model_record["model_name"],
fold_id=fold_id,
partition=partition,
metric=metric,
task_type=task_type,
scores=model_record["avg_scores"],
branch_id=model_record["branch_id"],
)
projected.flush()
run_result = RunResult(predictions=projected, per_dataset={spectro.name: {"engine": "dag-ml"}})
run_result._dagml_score_set = _cv_only_stacking_score_set(scores) # noqa: SLF001
run_result._dagml_refit_artifacts = [] # noqa: SLF001
return run_result
def _stacking_model_metadata(pipeline: list[Any]) -> dict[str, Any]:
"""Retain model-step controls when extracting the bare meta estimator."""
from nirs4all.pipeline.dagml_bridge import _step_to_dsl
model_steps = [step for step in pipeline if isinstance(step, dict) and "model" in step]
if not model_steps:
raise DagMlUnsupported("stacking needs a downstream meta-model step")
return dict(_step_to_dsl(model_steps[-1]).get("metadata") or {})
def _stacking_inner_cv(
task_type: str,
spectro: Any,
pool: list[int],
folds: list[tuple[list[int], list[int]]],
identity: Any,
random_state: int | None,
group_by_sample: dict[int, str] | None = None,
) -> dict[str, Any]:
"""Declare an identity-keyed nested splitter suitable for the target domain."""
if group_by_sample:
for training, _validation in folds:
if len({group_by_sample[int(sample)] for sample in training}) < 2:
raise DagMlUnsupported("grouped stacking requires at least two groups in every outer-training fold")
return {"kind": "group_kfold", "n_splits": 2}
if task_type != "classification":
return {"kind": "kfold", "n_splits": 2, "shuffle": False, "seed": random_state}
targets = np.asarray(_dataset_y_rows(spectro, pool), dtype=float)
if targets.size != len(pool) or targets.reshape(len(pool), -1).shape[1] != 1:
raise DagMlUnsupported("classification stacking requires exactly one class target per training sample")
values = targets.reshape(-1)
if not np.all(np.isfinite(values)):
raise DagMlUnsupported("classification stacking requires finite class targets")
label_by_sample = {
int(sample): (0.0 if float(value) == 0 else float(value)).hex()
for sample, value in zip(pool, values, strict=True)
}
for fold_train, _fold_validation in folds:
counts: dict[str, int] = {}
for sample in fold_train:
label = label_by_sample[int(sample)]
counts[label] = counts.get(label, 0) + 1
if len(counts) < 2 or any(count < 2 for count in counts.values()):
raise DagMlUnsupported(
"classification stacking requires every outer-training fold to contain at least two samples per class for nested stratification"
)
return {
"kind": "stratified_kfold",
"n_splits": 2,
"shuffle": True,
"seed": random_state,
"strata": {identity.to_wire(int(sample)): label_by_sample[int(sample)] for sample in pool},
}
def _assemble_stacking_dsl(
pipeline: list[Any], branches: list[list[Any]], meta_learner: Any, spectro: Any,
identity: Any, pool: list[int], folds: list[tuple[list[int], list[int]]], envelope: dict[str, Any], *,
task_type: str, random_state: int | None, group_by_sample: dict[int, str] | None,
source_layout: dict[str, Any] | None = None,
prediction_aggregations: list[dict[str, Any]] | None = None,
selection_metric: str = "rmse",
) -> tuple[dict[str, Any], dict[str, Any], list[str]]:
"""Declare the same nested OOF graph for concrete runs and whole-stack HPO."""
partial_profile = source_layout is not None and source_layout.get("schema") == "nirs4all.source-stacking-layout.v4"
if not partial_profile and source_layout is not None and source_layout.get("missing_source_policy", "error") != "error" and task_type == "classification":
raise DagMlUnsupported("late-fusion missing_source_policy='zero_with_indicator' currently requires regression")
meta_metadata = _stacking_model_metadata(pipeline)
from nirs4all.operators.models.meta import CoverageStrategy, MetaModel, TestAggregation
meta_wrapper = next((step.get("model") for step in reversed(pipeline)
if isinstance(step, dict) and isinstance(step.get("model"), MetaModel)), None)
fold_aggregation = (
meta_wrapper.stacking_config.test_aggregation
if meta_wrapper is not None else TestAggregation.MEAN
)
if fold_aggregation in (TestAggregation.BEST_FOLD, TestAggregation.WEIGHTED_MEAN):
meta_metadata["stacking_test_aggregation"] = "best" if fold_aggregation == TestAggregation.BEST_FOLD else "weighted"
meta_metadata["stacking_test_metric"] = selection_metric
if meta_wrapper is not None and meta_wrapper.stacking_config.coverage_strategy in (
CoverageStrategy.DROP_INCOMPLETE, CoverageStrategy.IMPUTE_MEAN,
CoverageStrategy.IMPUTE_ZERO, CoverageStrategy.IMPUTE_FOLD_MEAN,
):
meta_metadata["stacking_oof_coverage_contract"] = {
"min_coverage_ratio": meta_wrapper.stacking_config.min_coverage_ratio,
}
if meta_wrapper is not None and meta_wrapper.stacking_config.coverage_strategy in (
CoverageStrategy.IMPUTE_MEAN, CoverageStrategy.IMPUTE_ZERO, CoverageStrategy.IMPUTE_FOLD_MEAN,
):
# Nested OOF is complete by construction. The core checks every input
# source against the requested IDs; a genuine hole is rejected until
# the requested imputation policy has native coverage and lineage.
meta_metadata["stacking_missing_prediction_policy"] = "complete_inner_oof_no_imputation"
if meta_metadata.get("nirs4all_finetune_params"):
raise DagMlUnsupported(
"meta-model HPO requires a native whole-stack nested search; "
"reusing a precomputed OOF matrix would leak its inner selection targets"
)
import dag_ml
from nirs4all.pipeline.dagml.cli_runner import data_bindings_for_nodes, split_invocation_for
from nirs4all.pipeline.dagml_bridge import (
_META_MODEL_CONTROLLER_ID,
_META_MODEL_REF,
_SOURCE_RECIPE_CONTROLLER_ID,
_json_safe_params,
_qualname,
)
outer_fold_set = build_fold_set(identity, folds, set_id="folds.stacking.outer")
outer_partition_mode = outer_fold_set.get("partition_mode")
refit_oof = {"stacking_refit_oof": "partitioned_inner_v1"} if outer_partition_mode == "resampled" else {}
refit_policy = "require_full_coverage"
# Canonical DSL: one duplication branch with N base sub-pipelines (each on the FULL data) + a
# merge_model meta-node. The meta-node carries the bare sklearn meta-learner (FQN + params) and the
# _META_MODEL_REF (so its dedicated manifest is not a generic model-kind catch-all) and binds to the
# meta-model controller via metadata.controller_id.
canonical_dsl: dict[str, Any] = {
"id": "nirs4all-stacking",
"inner_cv": _stacking_inner_cv(task_type, spectro, pool, folds, identity, random_state, group_by_sample),
"steps": [
{"kind": "branch", "mode": "duplication", "branches": [_canonical_branch(branch, index) for index, branch in enumerate(branches)]},
{
"kind": "merge_model",
"id": _META_NODE_ID,
"include_original_data": False,
"operator": {"class": _qualname(meta_learner), "ref": _META_MODEL_REF},
"params": _json_safe_params(meta_learner),
"metadata": {
**meta_metadata,
"nirs4all_use_proba": bool(task_type == "classification" and meta_wrapper is not None and meta_wrapper.use_proba),
"controller_id": _META_MODEL_CONTROLLER_ID,
"stacking_oof_execution": "nested_oof_v1",
"stacking_oof_refit_contract": {"policy": refit_policy},
**refit_oof,
},
**({"selectors": prediction_aggregations} if prediction_aggregations else {}),
},
],
}
if partial_profile:
from .source_missing import build_prediction_availability
assert source_layout is not None
if prediction_aggregations:
raise DagMlUnsupported("partial-source late fusion requires the exact ordered source producers without prediction aggregation")
canonical_dsl["metadata"] = {"prediction_availability": build_prediction_availability(
spectro, identity, pool, target_values=_dataset_y_rows(spectro, pool),
classification=task_type == "classification",
)}
availability = canonical_dsl["metadata"]["prediction_availability"]
if source_layout["target_policy"] == "complete" and not np.asarray(availability["target_validity_masks"]).all():
raise DagMlUnsupported("partial late-fusion labels require explicit target_policy='per_target'")
if source_layout["missing_source_policy"] == "error" and not all(all(mask) for mask in availability["source_presence"].values()):
raise DagMlUnsupported("late fusion requires complete modalities unless zero_with_indicator is explicit")
canonical_dsl["steps"][1]["metadata"]["prediction_availability_meta"] = True
canonical_dsl["steps"][1]["metadata"]["nirs4all_use_proba"] = False
canonical_dsl["steps"][1]["sources"] = [
step["id"] for branch in canonical_dsl["steps"][0]["branches"]
for step in branch["steps"] if step["kind"] == "model"
]
if prediction_aggregations:
canonical_branches = {branch["id"]: branch for branch in canonical_dsl["steps"][0]["branches"]}
for selector in prediction_aggregations:
if isinstance(selector.get("select"), dict) and (
"fold_candidates_top_k" in selector["select"]
or "diverse_fold_candidates" in selector["select"]
):
if selector.get("metric") == "val_score":
selector["metric"] = selection_metric
continue
requested_names = selector.get("select")
if not isinstance(requested_names, list):
continue
branch = canonical_branches[selector["branch"]]
models = [step for step in branch["steps"] if step["kind"] == "model"]
selected_ids = [
step["id"]
for name in requested_names
for step in models
if step["operator"]["class"].rsplit(".", 1)[-1] == name
]
# Legacy skips unknown names. Preserve that behavior while making
# an entirely unmatched selection an explicit unsupported shape.
if not selected_ids:
raise DagMlUnsupported(
f"merge branch {selector['branch']} has no models matching {requested_names}"
)
selector["select"] = {"models": list(dict.fromkeys(selected_ids))}
# A legacy explicit selection fixes both membership and column order.
# The native graph supports ordered `sources`; without it the host sees
# the selected prediction inputs in lexical producer-id order.
if all(selector.get("aggregate") is None and selector.get("branch") in canonical_branches
and (selector.get("select", "all") == "all" or
isinstance(selector.get("select"), dict) and set(selector["select"]) == {"models"})
for selector in prediction_aggregations):
ordered_sources: list[str] = []
for selector in prediction_aggregations:
branch_models = [step["id"] for step in canonical_branches[selector["branch"]]["steps"] if step["kind"] == "model"]
requested = selector.get("select", "all")
ordered_sources.extend(branch_models if requested == "all" else requested["models"])
canonical_dsl["steps"][1]["sources"] = list(dict.fromkeys(ordered_sources))
if fold_aggregation in (TestAggregation.BEST_FOLD, TestAggregation.WEIGHTED_MEAN):
for branch in canonical_dsl["steps"][0]["branches"]:
for step in branch["steps"]:
if step["kind"] == "model":
step["metadata"] = {
**step.get("metadata", {}),
"nirs4all_stack_fold_capture": True,
"nirs4all_stack_outer_fold_ids": [fold["fold_id"] for fold in outer_fold_set["folds"]],
}
probability_branches = {
selector["branch"] for selector in prediction_aggregations or []
if selector.get("aggregate") == "proba_mean"
or selector.get("metadata", {}).get("prediction_output") == "proba"
}
if task_type == "classification" and (partial_profile or probability_branches or (meta_wrapper is not None and meta_wrapper.use_proba)):
from .operator_routing import route_graph_node
first_source_ports: dict[str, str] = {}
for branch in canonical_dsl["steps"][0]["branches"]:
if not partial_profile and branch["id"] not in probability_branches and not (meta_wrapper is not None and meta_wrapper.use_proba):
continue
for source_step in branch["steps"]:
if source_step["kind"] != "model":
continue
if not callable(getattr(route_graph_node(source_step), "predict_proba", None)):
if partial_profile:
raise DagMlUnsupported("partial-source classification requires genuine full-class predict_proba from every branch")
continue
source_step["prediction_output_ports"] = ["proba"]
first_source_ports[source_step["id"]] = "proba"
if first_source_ports:
canonical_dsl["steps"][1]["source_ports"] = first_source_ports
if source_layout is not None:
# Put source bindings in the DSL BEFORE compilation/fingerprinting,
# not in a callback-only graph mutation invisible to native provenance.
for index, branch in enumerate(canonical_dsl["steps"][0]["branches"]):
source_model_ids = [step["id"] for step in branch["steps"] if step["kind"] == "model"]
if partial_profile and len(source_model_ids) != 1:
raise DagMlUnsupported("partial-source preprocessing requires one model owner per source")
for step in branch["steps"]:
if partial_profile and step["kind"] == "transform":
# These unbound nodes declare recipes only. The source
# model clones and fits them on its native fit intersection.
step["metadata"] = {**step.get("metadata", {}),
"controller_id": _SOURCE_RECIPE_CONTROLLER_ID,
"nirs4all_source_recipe": {
"schema_version": 1, "model_node_id": source_model_ids[0],
"source_name": source_layout["sources"][index]["source_name"],
}}
if step["kind"] == "model":
step["metadata"] = {**step.get("metadata", {}), "nirs4all_source_stacking": {
"layout_fingerprint": source_layout["fingerprint"], "source": source_layout["sources"][index],
**({"missing_source_policy": source_layout["missing_source_policy"]} if "missing_source_policy" in source_layout else {}),
**({"target_policy": source_layout["target_policy"]} if partial_profile else {}),
}}
if partial_profile:
step["metadata"]["prediction_availability_source"] = source_layout["sources"][index]["source_name"]
if source_layout.get("kind") == "typed_source_blocks":
step["metadata"]["source_index"] = index
step["metadata"]["source_name"] = source_layout["sources"][index]["source_name"]
canonical_dsl["steps"][1]["metadata"]["nirs4all_source_stacking"] = source_layout
_bind_experimental_unit_design(canonical_dsl, spectro, identity, pool)
manifests = controller_manifests(canonical_dsl)
canonical_dsl["root_seed"] = random_state if random_state is not None else 0
graph = dag_ml.compile_pipeline_dsl_artifact_with_controllers(canonical_dsl, manifests).graph.to_dict()
model_ids = [node["id"] for node in graph["nodes"] if node["kind"] == "model"]
base_model_ids = [model_id for model_id in model_ids if model_id != _META_NODE_ID]
if not base_model_ids:
raise DagMlUnsupported("stacking compile produced no base model node")
if _META_NODE_ID not in model_ids:
raise DagMlUnsupported("stacking compile produced no meta-model node")
# One data_binding per BASE model node (each binds its `x` to the full source). The meta-node has NO
# data binding: its features are the base branches' OOF, delivered as prediction_inputs (not data).
canonical_dsl["data_bindings"] = data_bindings_for_nodes(base_model_ids, envelope)
canonical_dsl["split_invocation"] = split_invocation_for(identity, folds, n_splits=len(folds))
if group_by_sample:
canonical_dsl["split_invocation"]["fold_set"]["sample_groups"] = {
identity.to_wire(int(sample)): group for sample, group in group_by_sample.items()
}
return canonical_dsl, graph, base_model_ids
def _run_stacking_branch(pipeline: list[Any], branches: list[list[Any]], meta_learner: Any, spectro: Any, dataset_arg: str, cli: str, venv_python: str, run_dir: Path, metric: str, task_type: str, dataset_pickle: str | None = None, config_name: str = "", random_state: int | None = None, source_layout: dict[str, Any] | None = None, refit: bool = True, prediction_aggregations: list[dict[str, Any]] | None = None, downstream_meta_steps: list[dict[str, Any]] | None = None, meta_per_branch: bool = False) -> RunResult:
"""Run a duplication branch + ``{"merge": "predictions"}`` + meta-model as ONE native dag-ml run (#10).
Lowers each inner sub-pipeline to a canonical duplication branch (``mode: "duplication"`` — each base
model node gets the FULL fold data view) + a ``merge_model`` meta-node carrying the meta-learner (its
FQN as ``operator.class`` so the node runner instantiates it) bound to ``controller:nirs4all.meta_model``
(which declares ``consumes_oof_predictions``, so dag-ml's planner permits the base→meta ``requires_oof``
edges). dag-ml runs ONE native CV+refit:
* for each outer fold, the native nested scheduler fits base models in inner folds restricted to
outer training samples, then fits the meta-model on those inner OOF predictions;
* separately fitted outer base models predict the outer validation samples, which the already
fitted outer meta-model predicts without reading their targets. Producer columns and sample
identities are aligned by the native plan, not positional Python joins.
The meta producer's cross-fold OOF average is ``cv_best_score`` — the stacking ensemble's CV score.
``best_rmse`` (final test) is also native: in REFIT dag-ml delivers each base producer's held-out
TEST prediction to the meta-node as a SEPARATE off-fold input keyed ``…oof:refit`` (partition Test,
``fold_id=None``), alongside the full Validation OOF the meta fits on. The refit meta-model predicts
the test set from those base TEST meta-features and emits a scored ``(test, fold_id=None)`` block.
LEAKAGE INVARIANT: outer validation samples never participate in their outer meta-model's inner
base fits or meta fit; the TEST meta-features come from the base
models' TEST predictions (the ``:refit`` off-fold input, phase-gated to REFIT), never their OOF/train,
and never enter the FIT_CV meta-features.
Named branches retain separate base/meta prediction views of this same native outcome, including
its captured REFIT artifacts. Resampled outer folds use a separately declared partitioned inner
OOF preparation for the meta REFIT; that evidence never enters outer CV scoring.
"""
_, splitter = _split_pipeline(pipeline)
if splitter is None:
raise DagMlUnsupported("engine='dag-ml' requires a cross-validator step (e.g. KFold) in the pipeline")
identity = mint_identity(spectro)
named_duplication = _uses_named_duplication_branch(pipeline)
# The handled shape rejects any exclude step, so the CV universe is the full train pool.
pool = spectro.index_column("sample", {"partition": "train"})
folds = _build_folds(splitter, spectro, pool, set())
outer_fold_set = build_fold_set(identity, folds, set_id="folds.stacking.outer")
outer_partition_mode = outer_fold_set.get("partition_mode")
refit_oof = {"stacking_refit_oof": "partitioned_inner_v1"} if outer_partition_mode == "resampled" else {}
group_by_sample = _split_group_grain(splitter, spectro, pool)
envelope = build_envelope(spectro, identity, sample_ints=pool, group_by_sample=group_by_sample)
canonical_dsl, graph, base_model_ids = _assemble_stacking_dsl(
pipeline[:-len(downstream_meta_steps)] if downstream_meta_steps else pipeline,
branches, meta_learner, spectro, identity, pool, folds, envelope,
task_type=task_type, random_state=random_state, group_by_sample=group_by_sample, source_layout=source_layout,
prediction_aggregations=prediction_aggregations,
selection_metric=metric,
)
stacking_source_ports: dict[str, dict[str, str]] = {
_META_NODE_ID: canonical_dsl["steps"][1].get("source_ports", {}),
}
final_meta_node_id = _META_NODE_ID
final_meta_learner = meta_learner
stacking_source_orders: dict[str, list[str]] = {}
if canonical_dsl["steps"][1].get("sources"):
stacking_source_orders[_META_NODE_ID] = canonical_dsl["steps"][1]["sources"]
meta_source_nodes: dict[str, list[str]] = {}
meta_labels: dict[str, str] = {}
if downstream_meta_steps:
import dag_ml
from nirs4all.operators.models.meta import TestAggregation
from nirs4all.pipeline.dagml_bridge import _META_MODEL_CONTROLLER_ID, _META_MODEL_REF, _json_safe_params, _qualname
source_nodes: dict[str, list[str]] = {}
source_steps: dict[str, dict[str, Any]] = {}
source_estimators: dict[str, Any] = {}
from .operator_routing import route_graph_node
for branch in canonical_dsl["steps"][0]["branches"]:
for source_step in branch["steps"]:
if source_step["kind"] == "model":
source_name = source_step["operator"]["class"].rsplit(".", 1)[-1]
source_nodes.setdefault(source_name, []).append(source_step["id"])
source_steps[source_step["id"]] = source_step
source_estimators[source_step["id"]] = route_graph_node(source_step)
first_meta_step = pipeline[-len(downstream_meta_steps) - 1]
first_meta_name = first_meta_step.get("name") or first_meta_step["model"].name
meta_labels[_META_NODE_ID] = first_meta_name
selected_first_sources = [
model_id
for selector in prediction_aggregations or []
for model_id in (
selector["select"].get("models", [])
if isinstance(selector.get("select"), dict)
else ([selector["model"]] if selector.get("model") in base_model_ids else [])
)
]
meta_source_nodes[_META_NODE_ID] = list(dict.fromkeys(selected_first_sources)) or base_model_ids
source_nodes[first_meta_name] = [_META_NODE_ID]
source_steps[_META_NODE_ID] = canonical_dsl["steps"][-1]
source_estimators[_META_NODE_ID] = meta_learner
previous_meta_name = first_meta_name
for level, step in enumerate(downstream_meta_steps, start=2):
next_aggregation = step["model"].stacking_config.test_aggregation
if next_aggregation in (TestAggregation.BEST_FOLD, TestAggregation.WEIGHTED_MEAN):
canonical_dsl["steps"][-1]["metadata"].update({
"stacking_fold_test_capture": True,
"nirs4all_stack_fold_capture": True,
"nirs4all_stack_outer_fold_ids": [fold["fold_id"] for fold in build_fold_set(identity, folds, set_id="folds.stacking.outer")["folds"]],
})
final_meta_learner = step["model"].model
final_meta_node_id = f"{_META_NODE_ID}.level{level}"
requested_names = list(dict.fromkeys(step["model"].source_models))
explicit_sources = requested_names != [previous_meta_name]
requested_sources = [node_id for name in requested_names for node_id in source_nodes[name]]
sources = requested_sources if explicit_sources else []
meta_source_nodes[final_meta_node_id] = requested_sources
if sources:
stacking_source_orders[final_meta_node_id] = sources
source_ports: dict[str, str] = {}
if task_type == "classification" and step["model"].use_proba:
for source_id in requested_sources:
if not callable(getattr(source_estimators[source_id], "predict_proba", None)):
continue # Legacy uses class labels when a classifier has no probability method.
source_step = source_steps[source_id]
source_step["prediction_output_ports"] = ["proba"]
source_ports[source_id] = "proba"
if source_ports:
stacking_source_ports[final_meta_node_id] = source_ports
canonical_dsl["steps"].append({
"kind": "merge_model",
"id": final_meta_node_id,
"include_original_data": False,
"operator": {"class": _qualname(final_meta_learner), "ref": _META_MODEL_REF},
"params": _json_safe_params(final_meta_learner),
**({"sources": sources} if sources else {}),
**({"source_ports": source_ports} if source_ports else {}),
"metadata": {
**_stacking_model_metadata([step]),
"nirs4all_use_proba": bool(task_type == "classification" and step["model"].use_proba),
"controller_id": _META_MODEL_CONTROLLER_ID,
"stacking_oof_execution": "nested_oof_v1",
"stacking_oof_refit_contract": {"policy": "require_full_coverage"},
**refit_oof,
**({
"stacking_test_aggregation": "best" if next_aggregation == TestAggregation.BEST_FOLD else "weighted",
"stacking_test_metric": metric,
} if next_aggregation in (TestAggregation.BEST_FOLD, TestAggregation.WEIGHTED_MEAN) else {}),
},
})
current_name = step.get("name") or step["model"].name
meta_labels[final_meta_node_id] = current_name
source_nodes[current_name] = [final_meta_node_id]
source_steps[final_meta_node_id] = canonical_dsl["steps"][-1]
source_estimators[final_meta_node_id] = final_meta_learner
previous_meta_name = current_name
if any(step["model"].stacking_config.test_aggregation in (TestAggregation.BEST_FOLD, TestAggregation.WEIGHTED_MEAN)
for step in downstream_meta_steps):
outer_fold_ids = [fold["fold_id"] for fold in build_fold_set(identity, folds, set_id="folds.stacking.outer")["folds"]]
for branch in canonical_dsl["steps"][0]["branches"]:
for base_step in branch["steps"]:
if base_step["kind"] == "model":
base_step["metadata"] = {
**base_step.get("metadata", {}),
"nirs4all_stack_fold_capture": True,
"nirs4all_stack_outer_fold_ids": outer_fold_ids,
}
for meta_step in canonical_dsl["steps"][1:-1]:
if meta_step["kind"] == "merge_model":
meta_step["metadata"].update({
"stacking_fold_test_capture": True,
"nirs4all_stack_fold_capture": True,
"nirs4all_stack_outer_fold_ids": outer_fold_ids,
})
canonical_dsl["root_seed"] = random_state if random_state is not None else 0
graph = dag_ml.compile_pipeline_dsl_artifact_with_controllers(
canonical_dsl, controller_manifests(),
).graph.to_dict()
outcome = run_cv_refit_bundle(
dsl=canonical_dsl, envelope=envelope, graph=graph, dataset_path=dataset_arg, workdir=run_dir, dagml_cli=cli, venv_python=venv_python, selection_metric=metric, dataset_pickle=dataset_pickle, dataset=spectro, random_state=random_state, refit=refit
)
if outcome["returncode"] != 0:
_raise_run_failure(outcome, "dag-ml stacking run failed")
from nirs4all.operators.models.meta import MetaModel, TestAggregation
fold_aggregation = next((
step["model"].stacking_config.test_aggregation
for step in pipeline
if isinstance(step, dict) and isinstance(step.get("model"), MetaModel)
), TestAggregation.MEAN)
if fold_aggregation in (TestAggregation.BEST_FOLD, TestAggregation.WEIGHTED_MEAN) and outcome["refit_artifacts"]:
import json
import dag_ml
from .native_results import _producer_node_from_artifact_id
from .node_runner import _DagmlSelectedFoldEstimator
reports = outcome["scores"].get("reports", [])
outer_fold_ids = [fold["fold_id"] for fold in build_fold_set(identity, folds, set_id="folds.stacking.outer")["folds"]]
for artifact in outcome["refit_artifacts"]:
producer = _producer_node_from_artifact_id(artifact.get("artifact_id"))
if producer not in base_model_ids:
continue
fold_estimators = artifact.get("fold_estimators")
if not fold_estimators:
raise DagMlUnsupported(f"stacking fold aggregation has no captured CV estimators for {producer}")
if any(fold_id not in fold_estimators for fold_id in outer_fold_ids):
raise DagMlUnsupported(f"stacking fold aggregation is missing an outer-fold estimator for {producer}")
fold_estimators = {fold_id: fold_estimators[fold_id] for fold_id in outer_fold_ids}
request = json.dumps({
"producer_node": producer,
"fold_ids": list(fold_estimators),
"metric": metric,
"reports": reports,
})
if fold_aggregation == TestAggregation.BEST_FOLD:
selected = json.loads(dag_ml.select_stacking_fold_json(request)) # type: ignore[attr-defined]
artifact["estimator"] = _DagmlSelectedFoldEstimator(fold_estimators, selected_fold=selected)
else:
weights = json.loads(dag_ml.stacking_fold_weights_json(request)) # type: ignore[attr-defined]
artifact["estimator"] = _DagmlSelectedFoldEstimator(
fold_estimators, weights=dict(zip(fold_estimators, weights, strict=True)),
)
if downstream_meta_steps and outcome["refit_artifacts"]:
import json
import dag_ml
from .native_results import _producer_node_from_artifact_id
for level, step in enumerate(downstream_meta_steps, start=2):
aggregation = step["model"].stacking_config.test_aggregation
if aggregation not in (TestAggregation.BEST_FOLD, TestAggregation.WEIGHTED_MEAN):
continue
outer_fold_ids = [fold["fold_id"] for fold in build_fold_set(identity, folds, set_id="folds.stacking.outer")["folds"]]
reports = outcome["scores"].get("reports", [])
previous_meta_node = _META_NODE_ID if level == 2 else f"{_META_NODE_ID}.level{level - 1}"
required_nodes = [*base_model_ids, _META_NODE_ID, *(f"{_META_NODE_ID}.level{stage}" for stage in range(2, level))]
for artifact in outcome["refit_artifacts"]:
producer = _producer_node_from_artifact_id(artifact.get("artifact_id"))
if producer not in required_nodes:
continue
captured = artifact.get("fold_estimators")
if not isinstance(captured, dict) or any(fold not in captured for fold in outer_fold_ids):
raise DagMlUnsupported(f"level-{level} stacking fold aggregation lacks paired CV estimators for {producer}")
artifact["fold_estimators"] = {fold: captured[fold] for fold in outer_fold_ids}
if producer != previous_meta_node:
continue
request = json.dumps({
"producer_node": producer, "fold_ids": outer_fold_ids,
"metric": metric, "reports": reports,
})
if aggregation == TestAggregation.BEST_FOLD:
artifact["fold_selection"] = {"selected_fold": json.loads(dag_ml.select_stacking_fold_json(request))} # type: ignore[attr-defined]
else:
weights = json.loads(dag_ml.stacking_fold_weights_json(request)) # type: ignore[attr-defined]
artifact["fold_selection"] = {"weights": dict(zip(outer_fold_ids, weights, strict=True))}
# List form exposes the ensemble; named form also exposes each base producer.
model_label = f"MetaModel_{type(final_meta_learner).__name__}"
result: RunResult
if named_duplication or meta_per_branch:
from .named_stacking import project_named_stacking
result = project_named_stacking(
outcome, branches=branches, branch_names=_duplication_branch_names(pipeline, len(branches)),
base_model_ids=base_model_ids, meta_node_id=_META_NODE_ID, meta_learner=meta_learner,
spectro=spectro, identity=identity, metric=metric, task_type=task_type, config_name=config_name,
pipeline=pipeline, random_state=random_state, meta_per_branch=meta_per_branch,
)
else:
referenced_meta_nodes = {
source for sources in meta_source_nodes.values() for source in sources
if source in meta_source_nodes
}
terminal_meta_nodes = [node for node in meta_source_nodes if node not in referenced_meta_nodes]
if len(terminal_meta_nodes) > 1:
import dag_ml
from .envelope import target_names
from .native_results import _producer_node_from_artifact_id
from .public_batch import DagMLBatchResult
from .run_backend import _metric_objective
def closure(node: str) -> set[str]:
return {node}.union(*(
closure(source) if source in meta_source_nodes else {source}
for source in meta_source_nodes[node]
))
children: list[RunResult] = []
for node in terminal_meta_nodes:
own_producers = closure(node)
own_scores = {
**outcome["scores"],
"reports": [report for report in outcome["scores"].get("reports", [])
if report.get("producer_node") in own_producers],
}
own_artifacts = [
artifact for artifact in outcome["refit_artifacts"]
if _producer_node_from_artifact_id(artifact.get("artifact_id")) in own_producers
]
child = _scores_to_run_result(
own_scores, spectro.name, meta_labels[node], metric, task_type,
producer=node, config_name=config_name, results=outcome["results"],
identity=identity, refit_artifacts=own_artifacts,
)
child._dagml_target_names = target_names(spectro) # noqa: SLF001
child._dagml_stacking_replay_producer = node # noqa: SLF001
child._dagml_stacking_independent_terminal = True # noqa: SLF001
for metadata in child.per_dataset.values():
metadata["producer_node"] = node
metadata["output_topology"] = "independent_metamodel"
children.append(child)
result = DagMLBatchResult(children)
result._dagml_score_set = outcome["scores"] # noqa: SLF001
result._dagml_refit_artifacts = outcome["refit_artifacts"] # noqa: SLF001
result._dagml_node_results = outcome["results"] # noqa: SLF001
result._dagml_stacking_replay_producer = "" # noqa: SLF001
decision = dag_ml.select_candidate(
{"id": "select:terminal_metamodel", "metric": {"name": metric, "objective": _metric_objective(metric)},
"evaluation_scope": "oof", "require_finite": True},
[{"candidate_id": node, "metrics": {metric: child.cv_best_score}}
for node, child in zip(terminal_meta_nodes, children, strict=True)],
)
result._dagml_selection_decision = decision # noqa: SLF001
result._dagml_selected_run = children[terminal_meta_nodes.index(decision["selected_candidate_id"])] # noqa: SLF001
else:
result = _scores_to_run_result(
outcome["scores"], spectro.name, model_label, metric, task_type,
producer=final_meta_node_id, config_name=config_name, results=outcome["results"],
identity=identity, refit_artifacts=outcome["refit_artifacts"],
)
if source_layout is not None:
for view in [result, *getattr(result, "runs", [])]:
for metadata in view.per_dataset.values():
metadata["source_stacking"] = {"execution": "native_nested_oof_v1", "layout": source_layout,
"legacy_spectral_concatenation": False}
for row in view.predictions._buffer:
row["result_metadata"] = {**(row.get("result_metadata") or {}), "source_stacking_layout": source_layout}
if refit_oof:
evidence = {
"schema": "nirs4all.nested-stacking-evaluation.v1",
"outer_partition_mode": outer_partition_mode,
"outer_validation_occurrences": sum(len(validation) for _, validation in folds),
"outer_validation_sample_count": len({sample for _, validation in folds for sample in validation}),
"training_sample_count": len(pool),
"refit_oof": "partitioned_inner_v1",
"refit_oof_is_selection_evidence": False,
}
for view in [result, *getattr(result, "runs", [])]:
for metadata in view.per_dataset.values():
metadata["stacking_evaluation"] = evidence
for view in [result, *getattr(result, "runs", [])]:
view._dagml_stacking_selectors = ( # noqa: SLF001
prediction_aggregations
if getattr(view, "_dagml_stacking_replay_producer", _META_NODE_ID) == _META_NODE_ID
else None
)
view._dagml_stacking_outer_fold_ids = [fold["fold_id"] for fold in outer_fold_set["folds"]] # noqa: SLF001
view._dagml_stacking_producer_classes = {
step["id"]: step["operator"]["class"].rsplit(".", 1)[-1]
for branch in canonical_dsl["steps"][0]["branches"]
for step in branch["steps"] if step.get("kind") == "model"
} # noqa: SLF001
view._dagml_stacking_probability_producers = {
step["id"] for step in canonical_dsl["steps"]
if step.get("kind") == "merge_model"
and step.get("metadata", {}).get("nirs4all_prediction_output") == "proba"
} | {
step["id"] for branch in canonical_dsl["steps"][0]["branches"]
for step in branch["steps"]
if step.get("kind") == "model"
and step.get("metadata", {}).get("nirs4all_prediction_output") == "proba"
} # noqa: SLF001
view._dagml_stacking_source_orders = stacking_source_orders # noqa: SLF001
view._dagml_stacking_source_ports = stacking_source_ports # noqa: SLF001
return result