Source code for thesis.workflows.mrtrix3.nodes.fod

"""Response function, FOD estimation, and intensity normalisation nodes."""

from pathlib import Path

from nipype import Node
from nipype.interfaces.mrtrix3 import EstimateFOD, MTNormalise, ResponseSD




def prepare_response_node(
    out_dir: Path,
    algorithm: str = "dhollander",
    name: str = "response",
) -> Node:
    """Configure ``dwi2response`` (multi-tissue response estimation).

    The DWI ``in_file`` and ``in_mask`` inputs are wired by the workflow.
    Output filenames are pinned so the artefacts land predictably in
    ``out_dir``.

    Args:
        out_dir: Output directory for response text files.
        algorithm: dwi2response backend (``dhollander``, ``tournier``,
            ``msmt_5tt``, ``fa``, or ``manual``).
        name: Nipype node name.

    Returns:
        Configured Nipype Node.
    """
    out_dir.mkdir(parents=True, exist_ok=True)
    node = Node(ResponseSD(), name=name)
    node.inputs.algorithm = algorithm
    node.inputs.wm_file = str(out_dir / "response_wm.txt")
    if algorithm in ("dhollander", "msmt_5tt"):
        node.inputs.gm_file = str(out_dir / "response_gm.txt")
        node.inputs.csf_file = str(out_dir / "response_csf.txt")
    # dwi2response wraps mrconvert internally — pass force via args so a
    # restart can overwrite the previous run's response.txt files.
    node.inputs.args = "-force"
    return node





def prepare_fod_node(
    out_dir: Path,
    algorithm: str = "msmt_csd",
    name: str = "fod",
) -> Node:
    """Configure ``dwi2fod`` (constrained spherical deconvolution).

    Args:
        out_dir: Output directory for FOD images.
        algorithm: ``msmt_csd`` (multi-shell, multi-tissue) or ``csd``
            (single-shell, single-tissue).
        name: Nipype node name.

    Returns:
        Configured Nipype Node.
    """
    out_dir.mkdir(parents=True, exist_ok=True)
    node = Node(EstimateFOD(), name=name)
    node.inputs.algorithm = algorithm
    node.inputs.wm_odf = str(out_dir / "wm_fod.mif")
    if algorithm == "msmt_csd":
        node.inputs.gm_odf = str(out_dir / "gm.mif")
        node.inputs.csf_odf = str(out_dir / "csf.mif")
        # msmt_csd requires one -lmax per tissue (WM/GM/CSF). Nipype's
        # EstimateFOD defaults to [8] which expands to a single -lmax 8
        # flag — dwi2fod then errors with
        # "Number of lmaxes specified (1) does not match number of tissues (3)".
        # WM gets the standard lmax 8; GM and CSF are modelled as isotropic
        # (lmax 0) per the MRtrix MSMT-CSD recommendation.
        node.inputs.max_sh = [8, 0, 0]
    else:
        node.inputs.max_sh = [8]
    node.inputs.args = "-force"
    return node





def prepare_mtnormalise_node(
    out_dir: Path,
    name: str = "mtnormalise",
) -> Node:
    """Configure ``mtnormalise`` (intensity bias correction across tissues).

    Always produces normalised WM, GM, and CSF volumes; the workflow only
    builds this node for the multi-tissue (``msmt_csd``) path.

    Args:
        out_dir: Output directory for normalised FODs.
        name: Nipype node name.

    Returns:
        Configured Nipype Node.
    """
    out_dir.mkdir(parents=True, exist_ok=True)
    node = Node(MTNormalise(), name=name)
    node.inputs.out_file_wm = str(out_dir / "wm_fod_norm.mif")
    node.inputs.out_file_gm = str(out_dir / "gm_norm.mif")
    node.inputs.out_file_csf = str(out_dir / "csf_norm.mif")
    node.inputs.args = "-force"
    return node