Skip to content

Commit

Permalink
pytests field pyvista
Browse files Browse the repository at this point in the history
  • Loading branch information
@samjrholt
samjrholt committed Nov 9, 2023
1 parent ad955a7 commit 1cb446d
Show file tree
Hide file tree
Showing 2 changed files with 178 additions and 20 deletions.
41 changes: 21 additions & 20 deletions discretisedfield/plotting/pyvista_field.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -278,12 +278,6 @@ def volume(
"""

if self.field.nvdim != 1:
raise RuntimeError(
"Only meshes with scalar dimensions can be plotted not"
f" {self.field.nvdim=}."
)

if plotter is None:
plot = pv.Plotter()
else:
Expand Down Expand Up @@ -350,12 +344,6 @@ def valid(self, plotter=None, multiplier=None, filename=None, **kwargs):
"""

if self.field.nvdim != 3:
raise RuntimeError(
"Only meshes with 3 vector dimensions can be plotted not"
f" {self.field.mesh.region.ndim=}."
)

if plotter is None:
plot = pv.Plotter()
else:
Expand Down Expand Up @@ -465,6 +453,12 @@ def contour(
if scalars is None and self.field.nvdim > 1:
scalars = self.field.vdims[-1]

if self.field.nvdim > 1:
if contour_scalars is None:
contour_kwargs["scalars"] = self.field.vdims[-1]
else:
contour_kwargs["scalars"] = contour_scalars

if plotter is None:
plot = pv.Plotter()
else:
Expand All @@ -483,10 +477,6 @@ def contour(
field_pv["valid"].astype(bool)
).cell_data_to_point_data()

contour_kwargs["scalars"] = (
self.field.vdims[-1] if contour_scalars is None else contour_scalars
)

plot.add_mesh(
field_pv.contour(isosurfaces=isosurfaces, **contour_kwargs),
smooth_shading=True,
Expand All @@ -510,8 +500,8 @@ def streamlines(
scalars=None,
color_field=None,
filename=None,
streamlines_kwargs={"max_time": 10, "n_points": 20},
tube_kwargs={"radius": 0.05},
streamlines_kwargs={},
tube_kwargs={},
**kwargs,
):
"""``pyvista`` streamline plot.
Expand Down Expand Up @@ -554,12 +544,14 @@ def streamlines(
streamlines_kwargs : dict, optional
Keyword arguments for the `pyvista.streamlines` function that generates the
streamline geometry from the vector field data.
streamline geometry from the vector field data. If not provided, the default
keys are ``max_time=10`` and ``n_points=20``.
tube_kwargs : dict, optional
Keyword arguments for the `pyvista.tube` function that creates
tubes around the streamlines.
tubes around the streamlines. If not provided, the default keys
are ``radius=0.05``.
**kwargs
Expand Down Expand Up @@ -604,8 +596,17 @@ def streamlines(
field_pv["valid"].astype(bool)
).cell_data_to_point_data()

if "max_time" not in streamlines_kwargs:
streamlines_kwargs["max_time"] = 10
if "n_points" not in streamlines_kwargs:
streamlines_kwargs["n_points"] = 20
if "radius" not in tube_kwargs:
tube_kwargs["radius"] = 0.05

streamlines = field_pv.streamlines("field", **streamlines_kwargs)

print(streamlines.point_data)

plot.add_mesh(
streamlines.tube(**tube_kwargs),
scalars=scalars,
Expand Down
157 changes: 157 additions & 0 deletions discretisedfield/tests/test_field.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,10 +10,12 @@
import matplotlib.pyplot as plt
import numpy as np
import pytest
import pyvista as pv
import scipy.fft as spfft
import xarray as xr

import discretisedfield as df
import discretisedfield.plotting.util as plot_util

from .test_mesh import html_re as mesh_html_re

Expand Down Expand Up @@ -3888,6 +3890,161 @@ def test_plot_large_sample():
field.k3d.vector()


def test_pyvista_vector(test_field):
pv.OFF_SCREEN = True
# Default
test_field.pyvista.vector()

# Color field
test_field.pyvista.vector(color_field=test_field.a)

# Colormap
test_field.pyvista.vector(cmap="hsv")

# Scalar
test_field.pyvista.vector(scalars=test_field.vdims[0])

# Change vector
test_field.pyvista.vector(vector=plot_util.cone())

# Scale
test_field.pyvista.vector(scale=1e9)

# Multiplier
test_field.pyvista.vector(multiplier=1e-6)

# glyph_kwargs
test_field.pyvista.vector(glyph_kwargs={"progress_bar": True})

# Plotter
plotter = pv.Plotter()
test_field.pyvista.vector(plotter=plotter)
plotter.show()

# Exceptions
with pytest.raises(RuntimeError):
test_field.a.pyvista.vector()
with pytest.raises(ValueError):
test_field.pyvista.vector(color_field=test_field)


def test_pyvista_scalar(test_field):
pv.OFF_SCREEN = True
# Default
test_field.a.pyvista.scalar()

# Colormap
test_field.a.pyvista.scalar(cmap="hsv")

# Multiplier
test_field.a.pyvista.scalar(multiplier=1e-6)

# 3D
test_field.pyvista.scalar(scalars=test_field.vdims[0])

# Plotter
plotter = pv.Plotter()
test_field.pyvista.scalar(plotter=plotter)
plotter.show()


def test_pyvista_volume(test_field):
pv.OFF_SCREEN = True
# Default
test_field.a.pyvista.volume()

# Colormap
test_field.a.pyvista.volume(cmap="hsv")

# Multiplier
test_field.a.pyvista.volume(multiplier=1e-6)

# 3D
test_field.pyvista.volume(scalars=test_field.vdims[0])

# Plotter
plotter = pv.Plotter()
test_field.pyvista.volume(plotter=plotter)
plotter.show()


def test_pyvista_valid(test_field):
pv.OFF_SCREEN = True
# Default
test_field.pyvista.valid()

# Colormap
test_field.pyvista.valid(cmap="hsv")

# Multiplier
test_field.pyvista.valid(multiplier=1e-6)

# Scalar field
test_field.a.pyvista.valid()

# Plotter
plotter = pv.Plotter()
test_field.pyvista.valid(plotter=plotter)
plotter.show()


def test_pyvista_contour(test_field):
pv.OFF_SCREEN = True
# Default
test_field.pyvista.contour()

# Colormap
test_field.pyvista.contour(cmap="hsv")

# Multiplier
test_field.pyvista.contour(multiplier=1e-6)

# Scalar field
test_field.b.pyvista.contour()

# Plotter
plotter = pv.Plotter()
test_field.pyvista.contour(plotter=plotter)
plotter.show()

# Isosurface
test_field.pyvista.contour(isosurfaces=5)
test_field.pyvista.contour(isosurfaces=[0, 1])

# Contour scalars
test_field.pyvista.contour(contour_scalars=test_field.vdims[0])

# Color field
test_field.pyvista.contour(color_field=test_field.a)

# contour_kwargs
test_field.pyvista.contour(contour_kwargs={"progress_bar": True})


def test_pyvista_streamlines(test_field):
pv.OFF_SCREEN = True
# Default
test_field.pyvista.streamlines()

# Colormap
test_field.pyvista.streamlines(cmap="hsv")

# Multiplier
test_field.pyvista.streamlines(multiplier=1e3)

# Plotter
plotter = pv.Plotter()
test_field.pyvista.streamlines(plotter=plotter)
plotter.show()

# Color field
test_field.pyvista.streamlines(color_field=test_field.a)

# kwargs
test_field.pyvista.streamlines(streamlines_kwargs={"progress_bar": True})
test_field.pyvista.streamlines(tube_kwargs={"progress_bar": True})


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


Expand Down

0 comments on commit 1cb446d

Please sign in to comment.