Spatial

class Spatial(SpatialTransform):
    r"""Apply resampling, affine motion, and elastic deformation together.

    This transform can:

    1. resample to a new space,
    2. apply a global affine mapping, and
    3. apply a dense elastic field,

    using a single sampling grid.

    The convenience wrappers [`Resample`][torchio.Resample],
    [`Affine`][torchio.Affine], and
    [`ElasticDeformation`][torchio.ElasticDeformation] expose subsets
    of these parameters with sensible defaults.

    Args:
        target: Output space.  Can be one of:

            - A scalar or 3-tuple of floats: output voxel spacing in mm.
              E.g., `1` for 1 mm isotropic, `(0.5, 0.5, 2.0)` for
              anisotropic.
            - A random spacing spec, sampled once per call at apply time: a
              2-tuple `(lo, hi)` (uniform, per axis), a 6-tuple
              `(lo1, hi1, lo2, hi2, lo3, hi3)` (per-axis ranges), a `Choice`,
              or a `torch.distributions.Distribution`. E.g.,
              `target=(1, 2)` or `target=(2, 4, 2, 4, 3, 6)`.
            - A `str`: either a path to an image file, or the name of
              an image in the subject (e.g., `"t1"`).
            - An [`Image`][torchio.Image] instance.
            - A `(spatial_shape, affine)` pair.
            - `None` (default): the output grid matches the input grid.
        scales: Scale factors $(s_1, s_2, s_3)$ for each axis.
            If a single value $x$ is given, all axes use $x$.
            If two values $(a, b)$ are given,
            $s_i \sim \mathcal{U}(a, b)$.
            If six values $(a_1, b_1, a_2, b_2, a_3, b_3)$ are given,
            $s_i \sim \mathcal{U}(a_i, b_i)$ independently.
            A `torch.distributions.Distribution` may also be passed.
            For example, `scales=0.5` halves the apparent object
            size (zoom out), and `scales=2` doubles it (zoom in).
        degrees: Euler rotation angles $(\theta_1, \theta_2, \theta_3)$
            in degrees, following the same value/range/distribution
            convention as *scales*.
        translation: Translation $(t_1, t_2, t_3)$ in mm, following
            the same convention.  The direction depends on the image
            orientation: in RAS+, `translation=(10, 0, 0)` shifts
            10 mm to the right.
        isotropic: If `True`, sample a single scale factor and
            reuse it for all three axes.  *scales* must then be a
            scalar or 2-value range.
        center: Pivot point for rotation and scaling.
            `"image"` (default) uses the image center;
            `"origin"` uses the world-coordinate origin.
        control_points: Optional pre-computed coarse displacement
            field with shape `(n_i, n_j, n_k, 3)` in mm.  If given,
            *num_control_points*, *max_displacement*, and
            *locked_borders* are ignored.
        num_control_points: Number of control points along each
            dimension of the coarse grid.  Can be a single `int`
            (isotropic) or a 3-tuple.  Minimum is 4.  Smaller values
            produce smoother deformations.
        max_displacement: Maximum displacement at each control point,
            in mm.  Follows the same value/range/distribution
            convention as *scales*.  Zero (default) disables elastic
            deformation.
        locked_borders: Number of outer control-point layers whose
            displacement is forced to zero.  `0` keeps all
            displacements; `1` zeros the outermost layer; `2`
            (default) zeros the two outermost layers.
        affine_first: If `True` (default), apply the affine mapping
            before the elastic field.  If `False`, apply the elastic
            field first.  The difference is significant for large
            transforms.
        image_interpolation: `"linear"` (default) or `"nearest"`.
            Used for [`ScalarImage`][torchio.ScalarImage] instances.
        label_interpolation: `"nearest"` (default) or `"linear"`.
            Used for [`LabelMap`][torchio.LabelMap] instances.
        antialias: If `True`, apply Gaussian smoothing before
            downsampling intensity images.  Label maps are never
            smoothed.  The standard deviations follow
            [Cardoso et al., MICCAI 2015](https://link.springer.com/chapter/10.1007/978-3-319-24571-3_81).
        default_pad_value: Fill rule for out-of-bounds intensity
            voxels.  `"minimum"` (default), `"mean"`, `"otsu"`,
            or a numeric value.
        default_pad_label: Numeric fill value for out-of-bounds label
            voxels.
        **kwargs: See [`Transform`][torchio.Transform].

    Note:
        The randomizable parameters (`scales`, `degrees`, `translation`,
        `max_displacement`, and the spacing form of `target`) follow a common
        value/range/distribution convention: a scalar is deterministic, a
        2-tuple $(a, b)$ samples uniformly, a 3-tuple sets per-axis values, a
        6-tuple sets per-axis ranges, and a `Choice` or
        `torch.distributions.Distribution` samples from a discrete set or a
        distribution. Structural parameters (`center`, interpolation, padding)
        are not randomizable.

    Examples:
        >>> import torchio as tio
        >>> # Resample to 1 mm isotropic with a random rotation
        >>> transform = tio.Spatial(
        ...     target=1,
        ...     degrees=(-10, 10),
        ...     translation=(-5, 5),
        ... )
        >>> # Elastic deformation only
        >>> transform = tio.Spatial(
        ...     max_displacement=7.5,
        ...     num_control_points=7,
        ... )
        >>> transformed = transform(subject)
    """

    def __init__(
        self,
        *,
        target: TypeTarget = None,
        scales: TypeParameterValue = 1.0,
        degrees: TypeParameterValue = 0.0,
        translation: TypeParameterValue = 0.0,
        isotropic: bool = False,
        center: TypeCenter = "image",
        control_points: TypeControlPoints | None = None,
        num_control_points: int | TypeThreeInts = 7,
        max_displacement: TypeParameterValue = 0.0,
        locked_borders: int = 2,
        affine_first: bool = True,
        image_interpolation: TypeInterpolation = "linear",
        label_interpolation: TypeInterpolation = "nearest",
        antialias: bool = False,
        default_pad_value: TypePadValue | float = "minimum",
        default_pad_label: int | float = 0,
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)
        self.target = target
        _validate_isotropic(scales, isotropic)
        self.scales = _to_positive_range(scales)
        self.degrees = _to_parameter_range(degrees)
        self.translation = _to_parameter_range(translation)
        self.isotropic = isotropic
        self.center = _parse_center(center)
        self.control_points = (
            _parse_control_points(control_points)
            if control_points is not None
            else None
        )
        self.num_control_points = _parse_num_control_points(num_control_points)
        self.max_displacement = _to_nonnegative_parameter_range(max_displacement)
        self.locked_borders = _parse_locked_borders(locked_borders)
        if self.locked_borders == 2 and 4 in self.num_control_points:
            msg = (
                "locked_borders=2 with 4 control points along any axis yields an"
                " identity elastic field"
            )
            raise ValueError(msg)
        self.affine_first = affine_first
        self.image_interpolation = _parse_interpolation(image_interpolation)
        self.label_interpolation = _parse_interpolation(label_interpolation)
        self.antialias = antialias
        self.default_pad_value = _parse_default_pad_value(default_pad_value)
        if not isinstance(default_pad_label, Number):
            msg = f"default_pad_label must be numeric, got {type(default_pad_label)}"
            raise TypeError(msg)
        self.default_pad_label = float(default_pad_label)

    def make_params(self, batch: SubjectsBatch) -> dict[str, Any]:
        """Sample random parameters and resolve the output space.

        Scales, degrees, translation, and control-point displacements are
        sampled once and applied identically to every sample and every
        image in the batch.

        Returns:
            Dict of serializable parameters for `apply_transform` and
            history replay.
        """
        images = self._get_images(batch)
        if not images:
            return {"selected_images": []}

        _, first_batch = next(iter(images.items()))
        first_shape = _get_spatial_shape(first_batch)
        first_affine = first_batch.affines[0]

        sampled_scales = _sample_scales(self.scales, self.isotropic)
        sampled_degrees = self.degrees.sample()
        sampled_translation = self.translation.sample()
        has_affine = _has_affine_component(
            sampled_scales,
            sampled_degrees,
            sampled_translation,
        )
        control_points, max_displacement = _resolve_control_points(
            self.control_points,
            self.num_control_points,
            self.max_displacement,
            self.locked_borders,
        )
        has_elastic = control_points is not None

        if has_affine or has_elastic:
            _check_shared_space(images, first_shape, first_affine)

        target_space = _resolve_target_space(
            self.target,
            batch,
            first_shape,
            first_affine,
        )
        forward_affine = None
        if has_affine:
            forward_affine = _build_forward_affine(
                scales=sampled_scales,
                degrees=sampled_degrees,
                translation=sampled_translation,
                center=self.center,
                shape=first_shape,
                affine=first_affine,
            )

        return {
            "selected_images": list(images.keys()),
            "target": _serialize_space(target_space),
            "original": _serialize_space((first_shape, first_affine)),
            "affine_matrix": _serialize_matrix(forward_affine),
            "control_points": _serialize_control_points(control_points),
            "max_displacement": list(max_displacement) if max_displacement else None,
            "affine_first": self.affine_first,
            "image_interpolation": self.image_interpolation,
            "label_interpolation": self.label_interpolation,
            "antialias": self.antialias,
            "default_pad_value": self.default_pad_value,
            "default_pad_label": self.default_pad_label,
        }

    def apply_transform(
        self,
        batch: SubjectsBatch,
        params: dict[str, Any],
    ) -> SubjectsBatch:
        """Apply the spatial mapping to every selected image in *batch*.

        The sampling grid is built once from the parameters produced by
        `make_params` and reused for all images and all batch samples.
        """
        selected_images = params.get("selected_images", [])
        if not selected_images:
            return batch

        target_space = _deserialize_space(params["target"])
        affine_matrix = _deserialize_matrix(params["affine_matrix"])
        control_points = _deserialize_control_points(params["control_points"])
        max_displacement = _deserialize_max_displacement(params["max_displacement"])

        _apply_spatial_to_batch(
            batch=batch,
            image_names=selected_images,
            target_space=target_space,
            affine_matrix=affine_matrix,
            control_points=control_points,
            max_displacement=max_displacement,
            affine_first=params["affine_first"],
            image_interpolation=params["image_interpolation"],
            label_interpolation=params["label_interpolation"],
            antialias=params.get("antialias", False),
            default_pad_value=params["default_pad_value"],
            default_pad_label=float(params["default_pad_label"]),
        )
        return batch

    @property
    def invertible(self) -> bool:
        """Whether this transform can be inverted."""
        return True

    def inverse(self, params: dict[str, Any]) -> _SpatialInverse:
        """Build the inverse transform from recorded parameters.

        The affine component is inverted exactly.  The elastic component
        is approximated by negating the sampled displacement field.  The
        `affine_first` flag is flipped so that the inverse operations
        run in the opposite order.

        Args:
            params: The parameter dict produced by `make_params`.

        Returns:
            A `_SpatialInverse` that resamples back to the original grid.
        """
        affine_matrix = _deserialize_matrix(params["affine_matrix"])
        inverse_affine = None
        if affine_matrix is not None:
            inverse_affine = np.linalg.inv(affine_matrix)

        control_points = _deserialize_control_points(params["control_points"])
        inverse_control_points = None
        if control_points is not None:
            inverse_control_points = -control_points
        original_space = _deserialize_space(params["original"])
        if original_space is None:
            msg = "Spatial inverse needs the original output space"
            raise RuntimeError(msg)

        return _SpatialInverse(
            target=original_space,
            affine_matrix=inverse_affine,
            control_points=inverse_control_points,
            affine_first=not params["affine_first"],
            image_interpolation=params["image_interpolation"],
            label_interpolation=params["label_interpolation"],
            default_pad_value=params["default_pad_value"],
            default_pad_label=float(params["default_pad_label"]),
            copy=False,
        )