Source code for blissoda.tomo.nxtomo.multiturn

from __future__ import annotations

from typing import Any

SEQUENCES = frozenset({"tomo:multitomo", "tomo:pcotomo"})




def matches(processor: Any, entry: Any) -> bool:
    return processor._sequence_name(entry) in SEQUENCES



def _proj(entry: Any) -> dict[str, Any]:
    return entry.get("technique", {}).get("proj", {})


def _require_proj(entry: Any) -> dict[str, Any]:
    proj_info = _proj(entry)
    missing = [name for name in ("proj_n", "nb_turns") if name not in proj_info]
    if missing:
        missing_fields = ", ".join(f"technique.proj.{name}" for name in missing)
        raise ValueError(f"Missing multiturn NXtomo metadata: {missing_fields}")
    return proj_info


def _turns(entry: Any) -> int:
    proj_info = _require_proj(entry)
    return max(int(proj_info["nb_turns"]), 1)


def _proj_n(entry: Any) -> int:
    proj_info = _require_proj(entry)
    return max(int(proj_info["proj_n"]), 1)




def expanded_subscan_plan(processor: Any, entry: Any) -> list[tuple[str | None, int]]:
    plan = processor._subscan_plan(entry)
    turns = _turns(entry)
    expanded = []
    for kind, count in plan:
        actual_count = int(count)
        if processor._is_projection_kind(kind):
            actual_count = _proj_n(entry) * turns
        expanded.append((kind, actual_count))
    return expanded





def synthetic_translation(processor: Any, entry: Any, alias_name: str):
    return processor._position_array(entry, alias_name)





def segment_specs(processor: Any, entry: Any, records: list[dict[str, Any]]):
    turns = _turns(entry)
    if turns <= 1:
        return [{"label": None, "pieces": records}]

    specs = []
    for turn_index in range(turns):
        pieces = []
        for record in records:
            if not processor._is_projection_kind(record["kind"]):
                pieces.append(record)
                continue
            frames_per_turn = _proj_n(entry)
            start = record["frame_start"] + turn_index * frames_per_turn
            stop = start + frames_per_turn
            pieces.append(processor._slice_record(record, start, stop))
        specs.append({"label": f"turn_{turn_index + 1:03d}", "pieces": pieces})
    return specs