Repository: Project-MONAI/MONAI
File: monai/losses/deform.py (Python)

Code before:
    IEEE TMI: Transactions on Medical Imaging. 2019. eprint arXiv:1809.05231.

    Adapted from:
        VoxelMorph (https://github.com/voxelmorph/voxelmorph)
    """

    def __init__(self, normalize: bool = False, reduction: LossReduction | str = LossReduction.MEAN) -> None:
        """
        Args:
            normalize:
                Whether to divide out spatial sizes in order to make the computation roughly
                invariant to image scale (i.e. vector field sampling resolution). Defaults to False.
            reduction: {``"none"``, ``"mean"``, ``"sum"``}
                Specifies the reduction to apply to the output. Defaults to ``"mean"``.

                - ``"none"``: no reduction will be applied.
                - ``"mean"``: the sum of the output will be divided by the number of elements in the output.
                - ``"sum"``: the output will be summed.
        """
        super().__init__(reduction=LossReduction(reduction).value)
        self.normalize = normalize

    def forward(self, pred: torch.Tensor) -> torch.Tensor:
        """
        Args:
            pred: the shape should be BCH(WD)

        Raises:
            ValueError: When ``self.reduction`` is not one of ["mean", "sum", "none"].

        """
        if pred.ndim not in [3, 4, 5]:
            raise ValueError(f"Expecting 3-d, 4-d or 5-d pred, instead got pred of shape {pred.shape}")
        for i in range(pred.ndim - 2):
            if pred.shape[-i - 1] <= 2:
                raise ValueError(f"All spatial dimensions must be > 2, got spatial dimensions {pred.shape[2:]}")
        if pred.shape[1] != pred.ndim - 2:
            raise ValueError(
                f"Number of vector components, {pred.shape[1]}, does not match number of spatial dimensions, {pred.ndim - 2}"
            )

        # first order gradient
        first_order_gradient = [spatial_gradient(pred, dim) for dim in range(2, pred.ndim)]

        # spatial dimensions in a shape suited for broadcasting below
        if self.normalize:
            spatial_dims = torch.tensor(pred.shape, device=pred.device)[2:].reshape((1, -1) + (pred.ndim - 2) * (1,))

        diffusion = torch.tensor(0)
        for dim_1, g in enumerate(first_order_gradient):
            dim_1 += 2

Review comment (suggestion):
It would be worth mentioning here that `pred` is a vector field with `C` being equal to the number of spatial dimensions.

Code after:
    IEEE TMI: Transactions on Medical Imaging. 2019. eprint arXiv:1809.05231.

    Adapted from:
        VoxelMorph (https://github.com/voxelmorph/voxelmorph)
    """

    def __init__(self, normalize: bool = False, reduction: LossReduction | str = LossReduction.MEAN) -> None:
        """
        Args:
            normalize:
                Whether to divide out spatial sizes in order to make the computation roughly
                invariant to image scale (i.e. vector field sampling resolution). Defaults to False.
            reduction: {``"none"``, ``"mean"``, ``"sum"``}
                Specifies the reduction to apply to the output. Defaults to ``"mean"``.

                - ``"none"``: no reduction will be applied.
                - ``"mean"``: the sum of the output will be divided by the number of elements in the output.
                - ``"sum"``: the output will be summed.
        """
        super().__init__(reduction=LossReduction(reduction).value)
        self.normalize = normalize

    def forward(self, pred: torch.Tensor) -> torch.Tensor:
        """
        Args:
            pred:
                Predicted dense displacement field (DDF) with shape BCH[WD],
                where C is the number of spatial dimensions.
                Note that diffusion loss can only be calculated
                when the sizes of the DDF along all spatial dimensions are greater than 2.

        Raises:
            ValueError: When ``self.reduction`` is not one of ["mean", "sum", "none"].
            ValueError: When ``pred`` is not 3-d, 4-d or 5-d.
            ValueError: When any spatial dimension of ``pred`` has size less than or equal to 2.
            ValueError: When the number of channels of ``pred`` does not match the number of spatial dimensions.

        """
        if pred.ndim not in [3, 4, 5]:
            raise ValueError(f"Expecting 3-d, 4-d or 5-d pred, instead got pred of shape {pred.shape}")
        for i in range(pred.ndim - 2):
            if pred.shape[-i - 1] <= 2:
                raise ValueError(f"All spatial dimensions must be > 2, got spatial dimensions {pred.shape[2:]}")
        if pred.shape[1] != pred.ndim - 2:
            raise ValueError(
                f"Number of vector components, i.e. number of channels of the input DDF, {pred.shape[1]}, "
                f"does not match number of spatial dimensions, {pred.ndim - 2}"
            )

        # first order gradient
        first_order_gradient = [spatial_gradient(pred, dim) for dim in range(2, pred.ndim)]