"""
.. _gradients_example:

Compute Gradients of a Field
~~~~~~~~~~~~~~~~~~~~~~~~~~~~

Estimate the gradient of a scalar or vector field in a data set.

The ordering for the output gradient tuple will be
{du/dx, du/dy, du/dz, dv/dx, dv/dy, dv/dz, dw/dx, dw/dy, dw/dz} for
an input array {u, v, w}.

Showing the :func:`pyvista.DataSetFilters.compute_derivative` filter.
"""

import numpy as np

# sphinx_gallery_thumbnail_number = 1
import pyvista as pv
from pyvista import examples

# A vtkStructuredGrid - but could be any mesh type
mesh = examples.download_carotid()
mesh
###############################################################################
# Now compute the gradients of the ``vectors`` vector field in the point data
# of that mesh. This is as simple as calling
# :func:`pyvista.DataSetFilters.compute_derivative`.
mesh_g = mesh.compute_derivative(scalars="vectors")
mesh_g["gradient"]

###############################################################################
# .. note:: You can also use :func:`pyvista.DataSetFilters.compute_derivative` for
#   computing other derivative based quantities, such as divergence, vorticity,
#   and Q-criterion. See function documentation for options.

###############################################################################
# ``mesh_g["gradient"]`` is an ``N`` by 9 NumPy array of the gradients, so we
# could make a dictionary of NumPy arrays of the gradients like:


def gradients_to_dict(arr):
    """A helper method to label the gradients into a dictionary."""
    keys = np.array(
        ["du/dx", "du/dy", "du/dz", "dv/dx", "dv/dy", "dv/dz", "dw/dx", "dw/dy", "dw/dz"]
    )
    keys = keys.reshape((3, 3))[:, : arr.shape[1]].ravel()
    return dict(zip(keys, mesh_g["gradient"].T))


gradients = gradients_to_dict(mesh_g["gradient"])
gradients

###############################################################################
# And we can add all of those components as individual arrays back to the mesh
# by:
mesh_g.point_data.update(gradients)
mesh_g

###############################################################################

keys = np.array(list(gradients.keys())).reshape(3, 3)

p = pv.Plotter(shape=keys.shape)
for i in range(keys.shape[0]):
    for j in range(keys.shape[1]):
        name = keys[i, j]
        p.subplot(i, j)
        p.add_mesh(mesh_g.contour(scalars=name), scalars=name, opacity=0.75)
        p.add_mesh(mesh_g.outline(), color="k")
p.link_views()
p.view_isometric()
p.show()


###############################################################################
# And there you have it, the gradients for a vector field. We could also do
# this for a scalar  field like for the ``scalars`` field in the given dataset.
mesh_g = mesh.compute_derivative(scalars="scalars")

gradients = gradients_to_dict(mesh_g["gradient"])
gradients

###############################################################################

mesh_g.point_data.update(gradients)

keys = np.array(list(gradients.keys())).reshape(1, 3)

p = pv.Plotter(shape=keys.shape)

for i in range(keys.shape[0]):
    for j in range(keys.shape[1]):
        name = keys[i, j]
        p.subplot(i, j)
        p.add_mesh(mesh_g.contour(scalars=name), scalars=name, opacity=0.75)
        p.add_mesh(mesh_g.outline(), color="k")
p.link_views()
p.view_isometric()
p.show()