brainmint 0.1.0__py3-none-any.whl

brainmint/__init__.py

@@ -0,0 +1,12 @@
+"""BrainMint models, integrations, and inference tools for medical image synthesis."""
+
+from importlib.metadata import PackageNotFoundError
+from importlib.metadata import version as _metadata_version
+
+try:
+    from ._version import __version__
+except (ImportError, ModuleNotFoundError):
+    try:
+        __version__ = _metadata_version("brainmint")
+    except PackageNotFoundError:
+        __version__ = "0.0.0"

brainmint/_version.py

@@ -0,0 +1,24 @@
+# file generated by vcs-versioning
+# don't change, don't track in version control
+from __future__ import annotations
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+version: str
+__version__: str
+__version_tuple__: tuple[int | str, ...]
+version_tuple: tuple[int | str, ...]
+commit_id: str | None
+__commit_id__: str | None
+
+__version__ = version = '0.1.0'
+__version_tuple__ = version_tuple = (0, 1, 0)
+
+__commit_id__ = commit_id = None

brainmint/callbacks/log_images.py

@@ -0,0 +1,341 @@
+import gc
+import logging
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Sequence, Callable, Union
+
+import torch
+import contextlib
+import pytorch_lightning as pl
+from monai.transforms import SaveImage
+from monai.data import MetaTensor
+
+_LOG = logging.getLogger(__name__)
+
+
+def get_tb_writer(trainer):
+    logs = trainer.loggers if getattr(trainer, "loggers", None) else ([trainer.logger] if getattr(trainer, "logger", None) else [])
+    for lg in logs:
+        if isinstance(lg, pl.loggers.TensorBoardLogger):
+            return lg.experiment  # SummaryWriter
+    return None
+
+
+class SaveMRIImages(pl.Callback):
+    """
+    After each validation epoch:
+      • iterate the val dataloader,
+      • collect ONE sample per requested modality (e.g., T1w/T2w/T1ce/FLAIR),
+      • run module inference (preferred) or forward() as fallback,
+      • save paired NIfTI files into: <run_dir>/<subdir>/epoch-XXX/.
+
+    Filenames:
+      <tag>_<MODALITY>_input.nii.gz
+      <tag>_<MODALITY>_<OUTNAME>.nii.gz  (for each selected model output)
+    """
+
+    def __init__(
+        self,
+        modalities: Sequence[str] = ("T1w", "T2w", "T1ce", "FLAIR"),
+        tag: str = "val",
+        subdir: str = "val_samples",
+        dirpath: Optional[str] = None,  # Defaults to trainer.log_dir
+        infer_kwarg_keys: Sequence[str] = ("image",),  # Batch keys to pass to inference
+        infer_method_candidates: Sequence[str] = ("_run_inference",),
+        fallback_to_forward: bool = True,
+        output_names: Sequence[str] = ("output", "z_mu", "z_sigma"),  # Names for outputs, matching the order the model returns.
+        save_outputs: Sequence[Union[int, bool, str]] = (1, 0, 0),  # Output to save, Boolean Mask | Indices | Output Names
+        input_save_list: Optional[Sequence[str]] = None,
+        save_dtype: torch.dtype = torch.float32,
+        clamp_min: float = 0.0,
+        clamp_max: float = 1.0,
+        output_activation: Optional[str] = None,
+        separate_folder: bool = False,
+        dataset_module: Optional[Any] = None,
+        
+        
+    ) -> None:
+        super().__init__()
+        self.modalities = [m.lower() for m in modalities]
+        self.tag = tag
+        self.subdir = subdir
+        self.dirpath = dirpath
+        self.infer_method_candidates = list(infer_method_candidates)
+        self.fallback_to_forward = fallback_to_forward
+        self.infer_kwarg_keys = list(infer_kwarg_keys)
+        self.input_save_list = list(input_save_list) if input_save_list is not None else []
+
+        self.output_names = list(output_names)
+        self.save_indices = self._normalize_save_selector(save_outputs, self.output_names)
+
+        self.save_dtype = save_dtype if isinstance(save_dtype, torch.dtype) else getattr(
+            torch, str(save_dtype).strip().lower().removeprefix("torch.")
+        )
+        self.clamp_min = clamp_min
+        self.clamp_max = clamp_max
+        self.output_activation = output_activation.lower() if isinstance(output_activation, str) and output_activation else None
+        if self.output_activation not in (None, "sigmoid", "tanh"):
+            raise ValueError(f"SaveMRIImages: unsupported output_activation={output_activation}")
+        self.separate_folder = separate_folder
+        self.dataset_module = dataset_module
+
+    @staticmethod
+    def _normalize_save_selector(selector: Sequence[Union[int, bool, str]], names: Sequence[str]) -> List[int]:
+        """Normalize save_outputs → sorted unique indices."""
+        n = len(names)
+        lower_names = {nm.lower() for nm in names}
+        name_map = {nm.lower(): i for i, nm in enumerate(names)}
+
+        is_bool_mask = (n > 0 and len(selector) == n) and all(
+            isinstance(x, bool) or (isinstance(x, int) and x in (0, 1)) for x in selector
+        )
+        if is_bool_mask:
+            idxs = [i for i, flag in enumerate(selector) if bool(flag)]
+            return sorted(set(idxs))
+
+        is_indices = len(selector) <= n and all(isinstance(x, int) and 0 <= x < n for x in selector)
+        if is_indices:
+            return sorted(set(selector))
+
+        is_names = len(selector) <= n and all(isinstance(x, str) and x.lower() in lower_names for x in selector)
+        idxs: List[int] = []
+        if is_names:
+            for s in selector:
+                j = name_map.get(s.lower())
+                if j is not None:
+                    idxs.append(j)
+
+        if not idxs:
+            _LOG.error(f"SaveMRIImages Callback - Invalid output selector after filtering: {selector}")
+        return sorted(set(idxs))
+
+    def _resolve_infer_fn(self, pl_module: pl.LightningModule) -> Optional[Callable]:
+        for name in self.infer_method_candidates:
+            fn = getattr(pl_module, name, None)
+            if callable(fn):
+                _LOG.info(f"SaveMRIImages Callback - Using inference method '{name}'.")
+                return fn
+        if self.fallback_to_forward and hasattr(pl_module, "forward"):
+            _LOG.warning(
+                "SaveMRIImages Callback - None of %s found; falling back to forward().",
+                self.infer_method_candidates,
+            )
+            return pl_module.forward
+        _LOG.warning("SaveMRIImages Callback - No inference method and fallback disabled; will skip saving.")
+        return None
+
+    def _slice_like_batch(self, val: Any, batch_size: int, i: int):
+        """Return the i-th sample if batched along dim 0; else return as-is."""
+        if torch.is_tensor(val):
+            if val.dim() > 0 and val.size(0) == batch_size:
+                return val[i:i + 1]
+            return val
+        if isinstance(val, (list, tuple)) and len(val) == batch_size:
+            return val[i]
+        return val
+
+    def _build_infer_kwargs(
+        self,
+        pl_module: pl.LightningModule,
+        batch: Dict[str, Any],
+        batch_size: int,
+        i: int,
+    ) -> Dict[str, Any]:
+        """Build kwargs from configured keys; first entry is treated as input."""
+        if not self.infer_kwarg_keys:
+            return {}
+        kwargs: Dict[str, Any] = {}
+        for k in self.infer_kwarg_keys:
+            if k in batch:
+                v = self._slice_like_batch(batch[k], batch_size, i)
+                if torch.is_tensor(v):
+                    v = v.to(pl_module.device, non_blocking=True)
+                kwargs[k] = v
+            else:
+                raise KeyError(f"Key:{k} missing in batch, Available keys:{list(batch.keys())}")
+        return {kk: vv for kk, vv in kwargs.items() if vv is not None}
+
+    @torch.no_grad()
+    def on_validation_epoch_end(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None:
+        # Only the main/global process writes
+        if not trainer.is_global_zero:
+            return
+
+        run_dir = Path(self.dirpath) if self.dirpath else Path(trainer.log_dir or trainer.default_root_dir)
+        epoch_dir = run_dir / self.subdir / f"epoch-{int(pl_module.current_epoch):03d}"
+        epoch_dir.mkdir(parents=True, exist_ok=True)
+
+        infer_fn = self._resolve_infer_fn(pl_module)
+        if infer_fn is None:
+            return
+
+        # Choose which datamodule/dataloader to use
+        if self.dataset_module is not None:
+            dm = self.dataset_module
+            try:
+                if hasattr(dm, "setup"): dm.setup()
+            except Exception as e:
+                _LOG.warning(f"... prepare/setup custom dataset module: {e}")
+            if dm is None:
+                _LOG.warning("SaveMRIImages Callback - No datamodule; skipping.")
+                return
+            vloader = dm.val_dataloader()
+        else:
+            dm = trainer.datamodule
+            if dm is None:
+                _LOG.warning("SaveMRIImages Callback - No datamodule; skipping.")
+                return
+            vloader = dm.val_dataloader()
+        
+        want = set(self.modalities)
+        got: set[str] = set()
+        _LOG.info(
+            f"SaveMRIImages Callback - Collecting modalities {sorted(want)} for epoch {int(pl_module.current_epoch)}"
+        )
+
+        required = set(self.infer_kwarg_keys)
+
+        for batch in vloader:
+            if want == got:
+                break
+
+            missing = required - set(batch.keys())
+            if missing:
+                raise KeyError(
+                    f"SaveMRIImages: missing required keys: {', '.join(sorted(missing))}. Present: {', '.join(sorted(batch.keys()))}"
+                )
+
+            mods = batch.get("modality", None)
+            batch_size = int(batch[self.infer_kwarg_keys[0]].shape[0])  # Shape of First kwarg
+            missing_modality = mods is None
+            if missing_modality:
+                _LOG.warning(
+                    "SaveMRIImages Callback - Batch missing 'modality'; saving first available samples."
+                )
+                mod_list = [None] * batch_size
+            else:
+                mod_list = [m.lower() for m in batch["modality"]]
+
+            remaining = len(want - got)
+            saved_without_modality = 0
+            for i in range(batch_size):
+                if want == got:
+                    break
+
+                mod_i = mod_list[i]
+                if mod_i is not None:
+                    mod_i = mod_i.lower()
+                    if (mod_i not in want) or (mod_i in got):
+                        continue
+                    mod_label = mod_i
+                else:
+                    if saved_without_modality >= remaining:
+                        break
+                    mod_label = f"sample{i:03d}"
+                
+                infer_kwargs = self._build_infer_kwargs(
+                    pl_module=pl_module,
+                    batch=batch,
+                    batch_size=batch_size,
+                    i=i,
+                )
+
+                plugin = getattr(trainer.strategy, "precision_plugin", None)
+                ctx = plugin.forward_context() if plugin else contextlib.nullcontext()
+                try:
+                    with ctx, torch.no_grad():
+                        result = infer_fn(infer_kwargs)
+                except Exception as e:
+                    kw_types = {k: type(v).__name__ for k, v in infer_kwargs.items()}
+                    raise RuntimeError(
+                        f"SaveMRIImages: inference failed for modality '{mod_i}'. "
+                        f"Kwarg types: {kw_types}"
+                    ) from e
+
+                if isinstance(result, dict):
+                    outputs = [result.get(nm, None) for nm in self.output_names]
+                elif isinstance(result, (tuple, list)):
+                    outputs = list(result)
+                else:  # single tensor
+                    outputs = [result]
+
+                # Prepare & save input once
+                save_inputs = []
+                if not self.input_save_list:
+                    save_inputs.append(self.infer_kwarg_keys[0])  # Fist Key (infer_kwarg_keys) is the input
+                else:
+                    save_inputs = self.input_save_list
+                
+                for inp_name in save_inputs:
+                    inp = infer_kwargs[inp_name]
+                    if inp.shape[0] != 1:
+                        raise ValueError(f"SaveMRIImages - Input dimension error, Length of 1st Dim should be 1. Got:{inp.shape} ")
+                    inp_vol = inp[0].detach().to(dtype=self.save_dtype).clamp(self.clamp_min, self.clamp_max).cpu()
+                    if not isinstance(inp_vol, MetaTensor):
+                        inp_vol = MetaTensor(inp_vol, meta={"filename_or_obj": f"_.nii.gz"})
+                    input_saver = SaveImage(
+                        output_dir=str(epoch_dir),
+                        output_postfix=f"{self.tag}_{mod_label}_input_{inp_name}",
+                        output_ext=".nii.gz",
+                        separate_folder=self.separate_folder,
+                    )
+                    input_saver(inp_vol)
+
+                    del input_saver, inp_vol
+                    gc.collect()
+
+                for j in self.save_indices:
+                    if j >= len(outputs) or outputs[j] is None:
+                        continue
+                    out = outputs[j]
+                    out_vol = out[0].detach()
+                    if self.output_activation == "sigmoid":
+                        out_vol = torch.sigmoid(out_vol)
+                    elif self.output_activation == "tanh":
+                        out_vol = torch.tanh(out_vol)
+                    out_vol = out_vol.to(self.save_dtype).clamp(self.clamp_min, self.clamp_max).cpu()
+
+                    if isinstance(out_vol, MetaTensor):
+                        out_meta_vol = out_vol
+                    else:
+                        
+                        #TODO: FIX PASSING OF INPUT META INFO
+                        out_meta = {}#getattr(inp_vol, "meta", None) or {}
+                        out_meta_vol = MetaTensor(out_vol, meta=out_meta)
+
+                    out_name = self.output_names[j] if j < len(self.output_names) else f"out{j}"
+                    saver = SaveImage(
+                        output_dir=str(epoch_dir),
+                        output_postfix=f"{self.tag}_{mod_label}_output_{out_name}",
+                        output_ext=".nii.gz",
+                        separate_folder=self.separate_folder,
+                    )
+                    saver(out_meta_vol)
+
+                    del saver, out_meta_vol
+                    gc.collect()
+
+                if mod_i is None:
+                    saved_without_modality += 1
+                    if saved_without_modality >= remaining:
+                        got = set(want)
+                else:
+                    got.add(mod_i)
+                _LOG.info(
+                    f"SaveMRIImages Callback - Saved {mod_label} (input & {','.join(self.output_names[k] for k in self.save_indices if k < len(self.output_names))}) → {epoch_dir}"
+                )
+
+                del infer_kwargs, outputs
+                gc.collect()
+
+        missing = sorted(want - got)
+        if missing:
+            _LOG.warning(f"SaveMRIImages Callback - Missing modalities this epoch: {missing}")
+        else:
+            _LOG.info(f"SaveMRIImages Callback - Saved all requested modalities: {sorted(got)}")
+
+
+        # After saving, if we used a custom dataset, tear it down and free memory
+        if self.dataset_module is not None:
+            del vloader, dm
+            gc.collect()
+            torch.cuda.empty_cache()

brainmint/callbacks/modality_completion_schedule.py

@@ -0,0 +1,315 @@
+from __future__ import annotations
+
+import logging
+from typing import Any, Dict, Iterable, List, Mapping, Optional, Sequence
+
+import pytorch_lightning as pl
+
+from brainmint.utils.schedules import PiecewiseSchedule
+
+logger = logging.getLogger(__name__)
+
+
+def _normalize_probs(probs: Mapping[str, float]) -> Dict[str, float]:
+    out = {str(k): float(v) for k, v in dict(probs).items()}
+    s = float(sum(max(0.0, v) for v in out.values()))
+    if s <= 0.0:
+        raise ValueError(f"Cannot normalize probabilities with non-positive sum: {out}")
+    return {k: max(0.0, v) / s for k, v in out.items()}
+
+
+def _iter_transforms(root: Any) -> Iterable[Any]:
+    """Depth-first walk over transform graphs (matches tests/data/test_translation_configs.py)."""
+    if root is None:
+        return
+    stack = [root]
+    seen = set()
+    while stack:
+        current = stack.pop()
+        if id(current) in seen:
+            continue
+        seen.add(id(current))
+        yield current
+
+        children: List[Any] = []
+        nested = getattr(current, "_transform", None)
+        if nested is not None and nested is not current:
+            children.append(nested)
+
+        transforms = getattr(current, "transforms", None)
+        if transforms:
+            children.extend(list(transforms))
+
+        extra_start = getattr(current, "extra_xforms_start", None)
+        if extra_start:
+            children.extend(list(extra_start))
+
+        extra_end = getattr(current, "extra_xforms_end", None)
+        if extra_end:
+            children.extend(list(extra_end))
+
+        stack.extend(reversed(children))
+
+
+def _find_choice_states(dm: Any) -> List[Any]:
+    st: List[Any] = []
+    tf = getattr(dm, "train_tf", None)
+    root = getattr(tf, "transform", None) or tf
+    for t in _iter_transforms(root):
+        state = getattr(t, "state", None)
+        if state is not None and hasattr(state, "get_choices") and hasattr(state, "set_choices"):
+            st.append(state)
+    return st
+
+
+def _find_sampling_states(dm: Any) -> List[Any]:
+    st: List[Any] = []
+    tf = getattr(dm, "train_tf", None)
+    root = getattr(tf, "transform", None) or tf
+    for t in _iter_transforms(root):
+        ss = getattr(t, "sampling_state", None)
+        if ss is not None and hasattr(ss, "get_config") and hasattr(ss, "set_config"):
+            st.append(ss)
+    return st
+
+
+class ModalityCompletionScheduleCallback(pl.Callback):
+    """
+    Generic, config-driven epoch-boundary scheduler for modality completion experiments.
+
+    What it can update:
+      1) Bucket sampler probabilities via DataModule.set_bucket_probs(...)
+      2) Stream selection probabilities via SharedChoiceState.set_choices(...)
+      3) Partial sampling config via SharedSamplingState.set_config(...)
+
+    All behavior is driven by YAML config (step/linear schedules). No stage hardcoding.
+    """
+
+    def __init__(
+        self,
+        *,
+        bucket_schedule: Optional[Sequence[Mapping[str, Any]]] = None,
+        choice_schedule: Optional[Sequence[Mapping[str, Any]]] = None,
+        sampling_schedule: Optional[Sequence[Mapping[str, Any]]] = None,
+        update_on: str = "epoch_start",
+        strict: bool = True,
+        normalize_bucket_probs: bool = True,
+        normalize_choice_probs: bool = True,
+        auto_complement_two_way: bool = True,
+        log_updates: bool = True,
+    ) -> None:
+        super().__init__()
+        self.update_on = str(update_on).lower().strip()
+        if self.update_on not in ("epoch_start", "epoch_end"):
+            raise ValueError("update_on must be one of: 'epoch_start', 'epoch_end'")
+
+        self.strict = bool(strict)
+        self.normalize_bucket_probs = bool(normalize_bucket_probs)
+        self.normalize_choice_probs = bool(normalize_choice_probs)
+        self.auto_complement_two_way = bool(auto_complement_two_way)
+        self.log_updates = bool(log_updates)
+
+        self.bucket_sched = PiecewiseSchedule(bucket_schedule, name="bucket_schedule")
+        self.choice_sched = PiecewiseSchedule(choice_schedule, name="choice_schedule")
+        self.sampling_sched = PiecewiseSchedule(sampling_schedule, name="sampling_schedule")
+
+        self._choice_states: List[Any] = []
+        self._sampling_states: List[Any] = []
+
+    def setup(self, trainer: pl.Trainer, pl_module: pl.LightningModule, stage: Optional[str] = None) -> None:
+        if stage not in (None, "fit"):
+            return
+        dm = trainer.datamodule
+        self._choice_states = _find_choice_states(dm)
+        self._sampling_states = _find_sampling_states(dm)
+
+        if self.strict:
+            if (self.choice_sched.steps or self.choice_sched.lines) and not self._choice_states:
+                raise RuntimeError("choice_schedule provided but no SharedChoiceState found in train transforms.")
+            if (self.sampling_sched.steps or self.sampling_sched.lines) and not self._sampling_states:
+                raise RuntimeError("sampling_schedule provided but no SharedSamplingState found in train transforms.")
+
+    def on_train_epoch_start(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None:
+        if self.update_on == "epoch_start":
+            self._apply(trainer)
+
+    def on_train_epoch_end(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None:
+        if self.update_on == "epoch_end":
+            self._apply(trainer)
+
+    def _apply(self, trainer: pl.Trainer) -> None:
+        epoch = int(trainer.current_epoch)
+        dm = trainer.datamodule
+
+        bucket_probs = self.bucket_sched.value_at(epoch)
+        choice_updates = self.choice_sched.value_at(epoch)
+        sampling_updates = self.sampling_sched.value_at(epoch)
+
+        if bucket_probs is not None:
+            self._apply_bucket_probs(dm, bucket_probs, epoch)
+
+        if choice_updates is not None:
+            self._apply_choice_probs(choice_updates, epoch)
+
+        if sampling_updates is not None:
+            self._apply_sampling_config(sampling_updates)
+
+    # ---------------------------
+    # Bucket probs
+    # ---------------------------
+    def _apply_bucket_probs(self, dm: Any, value: Any, epoch: int) -> None:
+        if not isinstance(value, Mapping):
+            raise TypeError(f"bucket_schedule value must be a mapping bucket->prob, got {type(value)}")
+
+        probs = {str(k): float(v) for k, v in dict(value).items()}
+        if self.normalize_bucket_probs:
+            probs = _normalize_probs(probs)
+
+        setter = getattr(dm, "set_bucket_probs", None)
+        if not callable(setter):
+            msg = "DataModule has no set_bucket_probs(); cannot apply bucket_schedule."
+            if self.strict:
+                raise RuntimeError(msg)
+            logger.warning(msg)
+            return
+
+        setter(probs)
+
+        if self.log_updates:
+            logger.info("[Schedule] bucket_probs(epoch=%d): %s", int(epoch), probs)
+
+    # ---------------------------
+    # Choice probs (SharedChoiceState)
+    # ---------------------------
+    def _apply_choice_probs(self, value: Any, epoch: int) -> None:
+        if not isinstance(value, Mapping):
+            raise TypeError(f"choice_schedule value must be a nested mapping, got {type(value)}")
+
+        updates: Dict[str, Any] = dict(value)
+
+        for state in self._choice_states:
+            cur = dict(state.get_choices())
+            changed = False
+
+            for bucket, mods in updates.items():
+                bucket = str(bucket)
+                if bucket not in cur:
+                    if self.strict:
+                        raise KeyError(f"choice_schedule refers to unknown bucket '{bucket}'")
+                    logger.warning("choice_schedule: unknown bucket '%s' (ignored)", bucket)
+                    continue
+                if not isinstance(mods, Mapping):
+                    if self.strict:
+                        raise TypeError(f"choice_schedule[{bucket}] must be a mapping modality->alias_probs")
+                    continue
+
+                bcfg = cur[bucket]
+                if not isinstance(bcfg, Mapping):
+                    if self.strict:
+                        raise TypeError(f"choices[{bucket}] must be mapping")
+                    continue
+
+                for mod, alias_probs in mods.items():
+                    mod_key = str(mod).lower()
+                    if not isinstance(alias_probs, Mapping):
+                        if self.strict:
+                            raise TypeError(f"choice_schedule[{bucket}][{mod_key}] must be mapping alias->prob")
+                        continue
+
+                    # Find modality config; fall back to "*"
+                    mcfg = bcfg.get(mod_key, None)
+                    wildcard = False
+                    if mcfg is None:
+                        mcfg = bcfg.get("*", None)
+                        wildcard = True
+                    if mcfg is None:
+                        if self.strict:
+                            raise KeyError(f"choices missing config for bucket='{bucket}' modality='{mod_key}' (and no '*')")
+                        logger.warning("choices missing config for bucket='%s' modality='%s' (ignored)", bucket, mod_key)
+                        continue
+                    if not isinstance(mcfg, Mapping):
+                        if self.strict:
+                            raise TypeError(f"choices[{bucket}][{mod_key if not wildcard else '*'}] must be mapping")
+                        continue
+
+                    streams = dict(mcfg.get("streams") or {})
+                    if not streams:
+                        if self.strict:
+                            raise KeyError(f"choices[{bucket}][{mod_key if not wildcard else '*'}].streams is empty")
+                        continue
+
+                    probs = dict(mcfg.get("probs") or {})
+
+                    # apply updates (set)
+                    for alias, p in alias_probs.items():
+                        probs[str(alias)] = float(p)
+
+                    if self.strict:
+                        unknown_aliases = set(str(a) for a in alias_probs.keys()) - set(streams.keys())
+                        if unknown_aliases:
+                            raise KeyError(
+                                f"choice_schedule[{bucket}][{mod_key}] refers to unknown stream aliases: "
+                                f"{sorted(unknown_aliases)}"
+                            )
+
+                    # Optional complement for 2-way choices when only one alias was specified
+                    if self.auto_complement_two_way:
+                        defined_aliases = list(streams.keys())
+                        specified = [str(a) for a in alias_probs.keys() if str(a) in defined_aliases]
+                        if len(defined_aliases) == 2 and len(specified) == 1:
+                            a0 = specified[0]
+                            other = defined_aliases[0] if defined_aliases[1] == a0 else defined_aliases[1]
+                            p0 = float(probs.get(a0, 0.0))
+                            probs[other] = max(0.0, 1.0 - p0)
+
+                    # Keep only aliases defined in streams
+                    probs = {a: float(probs.get(a, 0.0)) for a in streams.keys()}
+
+                    if self.normalize_choice_probs:
+                        probs = _normalize_probs(probs)
+
+                    new_mcfg = dict(mcfg)
+                    new_mcfg["probs"] = probs
+
+                    new_bucket_cfg = dict(cur[bucket])
+                    if wildcard and "*" in new_bucket_cfg:
+                        new_bucket_cfg["*"] = new_mcfg
+                    else:
+                        new_bucket_cfg[mod_key] = new_mcfg
+                    cur[bucket] = new_bucket_cfg
+
+                    changed = True
+
+            if changed:
+                try:
+                    state.set_epoch(int(epoch))
+                except Exception:
+                    pass
+                state.set_choices(cur)
+                if self.log_updates:
+                    logger.info("[Schedule] choice_probs(epoch=%d): updated SharedChoiceState", epoch)
+
+    # ---------------------------
+    # Sampling config (SharedSamplingState)
+    # ---------------------------
+    def _apply_sampling_config(self, value: Any) -> None:
+        if not isinstance(value, Mapping):
+            raise TypeError(f"sampling_schedule value must be a mapping, got {type(value)}")
+        upd = {str(k): v for k, v in dict(value).items()}
+        if self.strict:
+            if "sigma_prob" in upd:
+                upd["sigma_prob"] = float(upd["sigma_prob"])
+                if not (0.0 <= upd["sigma_prob"] <= 1.0):
+                    raise ValueError("sampling_schedule sigma_prob must be between 0 and 1.")
+            if "sigma_alpha" in upd:
+                upd["sigma_alpha"] = float(upd["sigma_alpha"])
+                if upd["sigma_alpha"] < 0.0:
+                    raise ValueError("sampling_schedule sigma_alpha must be >= 0.")
+
+        for ss in self._sampling_states:
+            cur = dict(ss.get_config())
+            cur.update(upd)
+            ss.set_config(cur)
+
+        if self.log_updates:
+            logger.info("[Schedule] sampling_config: %s", upd)