Source code for thesis.workflows.hcp.operations.streamline_warping

"""Streamline (tractography output) warping to template space."""




def inverse_warp_streamlines_task(
    fdt_paths: str,
    warp_field: str,
    reference_image: str,
    output_dir: str,
    _ordering_signal: object = None,
) -> list:  # noqa: UP006
    """
    Warp ProbTrackX2 outputs from subject diffusion space to template space.

    Uses ANTs ApplyTransforms with the patient→template forward warp field
    directly (``invert_transform_flags=[False]``) and Linear interpolation.
    Handles ``fdt_paths.nii.gz``, ``seeds_to_*.nii.gz``, and
    ``waytotal.nii.gz``.

    All imports are performed inside the function body so that Nipype can
    serialise and execute it as a standalone Function node.

    Args:
        fdt_paths: Path to the directory containing ProbTrackX2 outputs.
        warp_field: Path to patient→template warp field.
        reference_image: Template-space reference image (defines output grid).
        output_dir: Output directory for warped files.
        _ordering_signal: Ordering-only input (ignored). Wired from
            ``probtrackx2.fdt_paths`` so this node runs only after ProbTrackX2
            has written its outputs into ``fdt_paths``; the directory itself
            comes from the outdir router, not from this value.

    Returns:
        List of paths to the warped output files.
    """
    import sys
    from pathlib import Path

    del _ordering_signal  # ordering dependency only; value intentionally unused

    from nipype.interfaces.ants import ApplyTransforms

    fdt_dir = Path(fdt_paths)
    out_dir = Path(output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    # Collect the recognised output files
    candidates = []
    fdt_file = fdt_dir / "fdt_paths.nii.gz"
    if fdt_file.exists():
        candidates.append(fdt_file)

    waytotal_file = fdt_dir / "waytotal.nii.gz"
    if waytotal_file.exists():
        candidates.append(waytotal_file)

    for seeds_file in sorted(fdt_dir.glob("seeds_to_*.nii.gz")):
        candidates.append(seeds_file)

    if not candidates:
        sys.stderr.write(f"WARNING: No recognised ProbTrackX2 output files found in {fdt_dir}\n")
        sys.stderr.flush()
        return []

    warped_files = []
    for src in candidates:
        out_file = str(out_dir / src.name)
        at = ApplyTransforms()
        at.inputs.input_image = str(src)
        at.inputs.reference_image = reference_image
        at.inputs.transforms = [warp_field]
        at.inputs.invert_transform_flags = [False]
        at.inputs.interpolation = "Linear"
        at.inputs.output_image = out_file
        sys.stderr.write(f"Warping {src.name} -> {out_file}\n")
        sys.stderr.flush()
        at.run()
        warped_files.append(out_file)

    return warped_files