# 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 `pyvista.DataSetFilters.compute_derivative()` filter.
import numpy as np

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
# `pyvista.DataSetFilters.compute_derivative()`.
mesh_g = mesh.compute_derivative(scalars='vectors')
mesh_g['gradient']

# NOTE:
#     You can also use `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):
    """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, strict=False))


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)

# `contour` makes the scalars it contours by the active scalars, so each mesh is
# colored by its own gradient component.
datasets = {name: mesh_g.contour(scalars=name) for name in keys.ravel()}

pv.plot_compare(
    datasets,
    opacity=0.75,
    reference_mesh=mesh_g.outline(),
    shape=keys.shape,
    cpos='iso',
)

# 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)

datasets = {name: mesh_g.contour(scalars=name) for name in keys.ravel()}

pv.plot_compare(
    datasets,
    opacity=0.75,
    reference_mesh=mesh_g.outline(),
    shape=keys.shape,
    cpos='iso',
)

# ----------------------------------------------------------------------
# Generated by sphinx-examples-as-code https://github.com/pyvista/sphinx-examples-as-code