VectorField

@beartype
class VectorField(AbstractField):
    """A vector field defined on a mesh.

    Each degree of freedom holds a 3-component vector value.

    """

    @property
    def dof_locations(self: Self) -> NDArray:
        """Return the coordinates of all vector degrees of freedom.

        Returns:
            Array of DOF coordinates with shape ``(n_dofs, 3)``.

        """
        return self.skfem_basis.doflocs.T[::3, :]

    @property
    def element_type(self: Self) -> skfem.Element:
        """Return the finite element type.

        Returns:
            A scikit-fem ``ElementVector`` wrapping the underlying scalar element.

        Notes:
            If the user has wrapped the scikit-fem element in
            ``ElementVector``, the element type is accepted as is.

        """
        if not isinstance(self._element_type, skfem.ElementVector):
            self._element_type = skfem.ElementVector(self._element_type)
        return self._element_type

    @property
    def value(self: Self) -> NDArray:
        """Return the field values at the degrees of freedom.

        Returns:
            A 2D array of vector values with shape ``(n_dofs, 3)``.

        """
        return self._value

    @value.setter
    def value(
        self: Self,
        value: Number | list | tuple | NDArray | Callable,
    ):
        """Set the field values.

        Args:
            value: New field values. Accepts a length-3 list or tuple (uniform
                vector), a 1D array of length 3 (uniform vector), a 2D array of
                shape ``(n_dofs, 3)``, or a callable that maps DOF coordinates
                to a ``(n_dofs, 3)`` array.

        Raises:
            TypeError: If a scalar is passed, or if a callable returns a
                non-array.
            ValueError: If the array shape is not compatible.

        """
        if isinstance(value, (int, float)):
            raise TypeError("Value of a vector field cannot be a number.")
        elif isinstance(value, (list, tuple)):
            if len(value) != 3:
                raise ValueError(
                    f"The length of value {type(value)} must be 3 for a vector "
                    f"field. Instead, length {len(value)} was received."
                )
            else:
                self._value = np.tile(np.array(value), (self.dof_locations.shape[0], 1))
        elif isinstance(value, np.ndarray):
            if value.shape == (3,):
                self._value = np.tile(value, (self.dof_locations.shape[0], 1))
            elif value.shape == (self.dof_locations.shape[0], 3):
                self._value = value
            else:
                raise ValueError(
                    "The value numpy array for the vector field must be a "
                    "1D array of length 3 or a 2D array of shape "
                    f" ({self.dof_locations.shape[0]}, 3). "
                    f"Instead, an array of shape {value.shape} was passed."
                )
        elif isinstance(value, Callable):
            out_arr = value(self.dof_locations)
            if not isinstance(out_arr, np.ndarray):
                raise TypeError(
                    "The output of the function used to set the value of the field "
                    "must be a numpy array. "
                    f"Instead, {type(out_arr).__name__} was obtained."
                )
            elif out_arr.shape != (self.dof_locations.shape[0], 3):
                raise ValueError(
                    "The output of the function used to set the value of the "
                    "vector field should have shape "
                    f"({self.dof_locations.shape[0]}, 3). "
                    f"Instead, {out_arr.shape} was received."
                )
            self._value = out_arr

    def set_subregion_value(
        self: Self, name: str, value: Number | list | tuple | NDArray | Callable
    ):
        """Set field values on a named subregion.

        Args:
            name: Name of the subregion to modify.
            value: New values for the subregion DOFs. Accepts the same types
                as the ``value`` setter, except scalars.

        Raises:
            RuntimeError: If no subregions are defined on the mesh.
            ValueError: Same as ``value`` setter or the subregion name does not exist.
            TypeError: Same as ``value`` setter.

        """
        if not self.mesh.subregions:
            raise RuntimeError("No subregion defined!")
        else:
            subregion_names = self.mesh.subregions.keys()
        if name not in subregion_names:
            raise ValueError(
                f"The subregion {name} does not exist. "
                f"Please select one from {subregion_names}."
            )

        dof_ids = self.get_subregion_dof_index(name)
        subregion_dofs = self.dof_locations[dof_ids]

        if isinstance(value, (int, float)):
            raise TypeError("Value of a vector field cannot be a number.")
        elif isinstance(value, (list, tuple)):
            if len(value) != 3:
                raise ValueError(
                    f"The length of value {type(value)} must be 3 for a vector "
                    f"field. Instead, length {len(value)} was received."
                )
            else:
                self._value[dof_ids] = np.tile(
                    np.array(value), (subregion_dofs.shape[0], 1)
                )
        elif isinstance(value, np.ndarray):
            if value.shape == (3,):
                self._value[dof_ids] = np.tile(value, (subregion_dofs.shape[0], 1))
            elif value.shape == (subregion_dofs.shape[0], 3):
                self._value[dof_ids] = value
            else:
                raise ValueError(
                    "The value numpy array for the vector field must be a "
                    "1D array of length 3 or a 2D array of shape "
                    f" ({subregion_dofs.shape[0]}, 3). "
                    f"Instead, an array of shape {value.shape} was passed."
                )
        elif isinstance(value, Callable):
            out_arr = value(subregion_dofs)
            if not isinstance(out_arr, np.ndarray):
                raise TypeError(
                    "The output of the function used to set the value of the field "
                    "must be a numpy array. "
                    f"Instead, {type(out_arr).__name__} was obtained."
                )
            elif out_arr.shape != (subregion_dofs.shape[0], 3):
                raise ValueError(
                    "The output of the function used to set the value of the "
                    "vector field should have shape "
                    f"({subregion_dofs.shape[0]}, 3). "
                    f"Instead, {out_arr.shape} was received."
                )
            self._value[dof_ids] = out_arr

    def to_skfem(self: Self) -> skfem.DiscreteField:
        """Convert the field to a scikit-fem discrete field.

        Returns:
            A scikit-fem ``DiscreteField`` interpolated from the field values.

        """
        return self.skfem_basis.interpolate(self.value.flatten("C"))

    def to_pyvista(self: Self) -> pv.UnstructuredGrid:
        """Convert the field to a PyVista UnstructuredGrid object.

        Returns:
            PyVista ``UnstructuredGrid`` with ``value`` vector data.

        Notes:
            Higher-order elements are projected to linear P1 elements for
            visualisation.

        """
        if isinstance(
            self.element_type.elem,
            (
                # Check for DOFs on the nodes or cells.
                skfem.ElementLineP1,
                skfem.ElementTriP1,
                skfem.ElementTetP1,
                skfem.ElementQuad1,
                skfem.ElementHex1,
            ),
        ):
            mesh_pv_main: pv.UnstructuredGrid = self.mesh.to_pyvista()

            if mesh_pv_main.points.shape == self.dof_locations.shape and np.allclose(
                mesh_pv_main.points, self.dof_locations
            ):
                final_value = self.value
            else:
                tree = KDTree(mesh_pv_main.points)
                dist, new_point_indices = tree.query(self.dof_locations)
                if not np.allclose(np.zeros_like(dist), dist):
                    raise RuntimeError(
                        "Unable to create new indices for scikit-fem mesh points."
                    )
                final_value = np.zeros_like(mesh_pv_main.points)
                final_value[new_point_indices] = self.value

            mesh_pv_main.point_data["value"] = final_value
            mesh_pv_main.set_active_vectors("value", "point")
            return mesh_pv_main
        elif isinstance(
            self.element_type.elem,
            (
                # Check for DOFs on the nodes or cells.
                skfem.ElementLineP0,
                skfem.ElementTriP0,
                skfem.ElementTetP0,
                skfem.ElementQuad0,
                skfem.ElementHex0,
            ),
        ):
            mesh_pv_main: pv.UnstructuredGrid = self.mesh.to_pyvista()

            if (
                mesh_pv_main.cell_centers().points.shape == self.dof_locations.shape
                and np.allclose(mesh_pv_main.cell_centers().points, self.dof_locations)
            ):
                final_value = self.value
            else:
                tree = KDTree(mesh_pv_main.cell_centers().points)
                dist, new_point_indices = tree.query(self.dof_locations)
                if not np.allclose(np.zeros_like(dist), dist):
                    raise RuntimeError(
                        "Unable to create new indices for dof_locations."
                    )
                final_value = np.zeros_like(mesh_pv_main.cell_centers().points)
                final_value[new_point_indices] = self.value
            mesh_pv_main.cell_data["value"] = final_value
            mesh_pv_main.set_active_vectors("value", "cell")
            return mesh_pv_main
        else:
            skfem_mesh = self.mesh.to_skfem()

            if isinstance(skfem_mesh, skfem.mesh.MeshLine1):
                new_element = skfem.ElementLineP1()
            elif isinstance(skfem_mesh, skfem.mesh.MeshTri1):
                new_element = skfem.ElementTriP1()
            elif isinstance(skfem_mesh, skfem.mesh.MeshTet1):
                new_element = skfem.ElementTetP1()
            elif isinstance(skfem_mesh, skfem.mesh.MeshQuad1):
                new_element = skfem.ElementQuad1()
            elif isinstance(skfem_mesh, skfem.mesh.MeshHex1):
                new_element = skfem.ElementHex1()
            else:
                raise TypeError(f"Unrecognised mesh type {skfem_mesh}.")

            flattened_value = self.skfem_basis.with_element(
                skfem.ElementVector(new_element)
            ).project(self.to_skfem())
            new_value = flattened_value[self.skfem_basis.nodal_dofs].T
            new_field = VectorField(self.mesh, new_element, new_value)
            return new_field.to_pyvista()

    def plot(
        self: Self,
        subregions: list[str] | None = None,
        colour_with: str = "z",
        factor: float = 1.0,
        cmap: str = "viridis",
        **kwargs,
    ):
        """Plot the vector field using PyVista glyphs.

        Args:
            subregions: Optional list defining subregions to plot.
            colour_with: Component used to colour the glyphs. One of ``"x"``,
                ``"y"``, ``"z"``, or ``"magnitude"``. Defaults to ``"z"``.
            factor: Scaling factor for glyphs. Defaults to 1.0.
            cmap: Matplotlib colormap name. Defaults to ``"viridis"``.
            **kwargs: Additional keyword arguments forwarded to
                ``pyvista.Plotter.show``.

        """
        if (
            subregions
            and self.mesh.subregions
            and not set(subregions).issubset(set(self.mesh.subregions))
        ):
            raise ValueError(
                f"Subregions to be plotted {subregions} "
                f"is not a subset of mesh subregions {self.mesh.subregions.keys()}."
            )
        elif subregions and self.mesh.subregions is None:
            raise ValueError(
                f"Cannot plot subregions {subregions} "
                "since no subregions are defined on the mesh."
            )
        plot_pv_obj: pv.UnstructuredGrid = self.to_pyvista()
        vector_value = plot_pv_obj["value"]
        magnitude_value = np.linalg.norm(vector_value, axis=-1)
        match colour_with.lower():
            case "x":
                colour_with_arr = vector_value[:, 0]
            case "y":
                colour_with_arr = vector_value[:, 1]
            case "z":
                colour_with_arr = vector_value[:, 2]
            case "magnitude":
                colour_with_arr = magnitude_value
            case _:
                raise ValueError(
                    "The value of colour_with must be one of x, y, z, or magnitude. "
                    f"Instead {colour_with} was passed."
                )
        label = (
            colour_with.lower()
            if colour_with.lower() == "magnitude"
            else f"{colour_with.lower()} component"
        )

        plot_pv_obj["magnitude"] = magnitude_value
        plot_pv_obj[label] = colour_with_arr
        plot_pv_obj.set_active_scalars(name=label)

        pl = pv.Plotter()
        if subregions:
            for subregion in subregions:
                pl.add_mesh(
                    plot_pv_obj.threshold(value=True, scalars=subregion).glyph(
                        scale="magnitude", factor=factor
                    ),
                    scalars=label,
                    cmap=cmap,
                    show_scalar_bar=True,
                )
        else:
            pl.add_mesh(
                plot_pv_obj.glyph(scale="magnitude", factor=factor),
                scalars=label,
                cmap=cmap,
                show_scalar_bar=True,
            )
        pl.show_axes_all()
        pl.show(**kwargs)

    def __str__(self: Self) -> str:
        """Return a concise string representation.

        Returns:
            Describes the field type, element type, subregions, and value range.

        """
        element_type = type(self.element_type.elem).__name__
        s = f"{self.__class__.__name__}(elements={element_type}"
        if self.mesh.subregions:
            s += f", subregions={[s for s in self.mesh.subregions]}"
        s += f", range=[{self.value.min()}, {self.value.max()}])"

        return s

    def __repr__(self: Self) -> str:
        """Return an unambiguous string representation.

        Returns:
            Same as ``__str__``.

        """
        return self.__str__()

    def __eq__(self: Self, other: object) -> bool:
        """Compare two vector fields for equality.

        Returns:
            ``True`` if same mesh, element type, and values (numpy tolerance).

        """
        if not isinstance(other, self.__class__):
            return False
        return (
            self.mesh == other.mesh
            and type(self.element_type.elem).__name__
            == type(other.element_type.elem).__name__
            and np.allclose(self.value, other.value)
        )

    def __call__(self: Self, point: list | tuple | NDArray) -> NDArray:
        """Evaluate the vector field at one or more points.

        Args:
            point: A single 3D point (length-3 list, tuple, or 1D array) or a
                2D array of points with shape ``(n_points, 3)``.

        Returns:
            Interpolated vector values at the points, with shape ``(n_points, 3)``.

        Raises:
            ValueError: If the point coordinates are not 3D or have an
                unsupported shape. A value error is raised from scikit-fem
                if the point is outside the mesh bounds.

        """
        if isinstance(point, list | tuple):
            if len(point) != 3:
                raise ValueError(
                    "The value of point for which to calculate the value of field "
                    f"needs to be a length 3 {type(point).__name__}. instead {len(point)=}."
                )
            else:
                point = np.array([point]).T
        elif isinstance(point, np.ndarray):
            if point.ndim == 1 and len(point) != 3:
                raise ValueError(
                    "The length of 1D point value array must be 3. "
                    f"instead {len(point)=}."
                )
            elif point.ndim == 2 and point.shape[-1] != 3:
                raise ValueError(
                    "The last dimension of 2D point value array must be 3. "
                    f"instead it is {point.shape[-1]}."
                )
            elif point.ndim > 2:
                raise ValueError(
                    "The point value array must be 1D of length 3 "
                    "or 2D with last dimension as 3. "
                    f"Instead {point.shape=}."
                )
            else:
                point = np.array([point]).T if point.ndim == 1 else point.T

        calc_fun = self.skfem_basis.interpolator(self.value.flatten(order="C"))
        return calc_fun(point).T

    def get_subregion_dof_index(self: Self, name: str) -> NDArray:
        """Return DOF indices belonging to a named subregion or field boundary.

        Args:
            name: Name of the subregion.

        Returns:
            Array of DOF indices, accounting for the three-component vector structure.

        Raises:
            ValueError: If the subregion is not found in either subdomains or
                boundaries.

        """
        skfem_elem_dofs = np.array(
            self.skfem_basis.element_dofs.T[:, ::3] / 3, dtype=np.int_
        )
        skfem_mesh = self.mesh.to_skfem()

        if name.lower() == "boundary":
            return np.array(
                self.skfem_basis.get_dofs().flatten()[::3] / 3, dtype=np.int_
            )
        if skfem_mesh.subdomains is not None and name in skfem_mesh.subdomains:
            dof_indices = np.unique(skfem_elem_dofs[skfem_mesh.subdomains[name]])
            return dof_indices
        elif skfem_mesh.boundaries is not None and name in skfem_mesh.boundaries:
            dof_indices = np.array(
                self.skfem_basis.get_dofs(name).flatten()[::3] / 3,
                dtype=np.int_,
            )
            return dof_indices
        else:
            raise ValueError(f"Unable to find subregion {name}.")