fix(plots): fix 3D plots in MP7/PRT examples (#181)

* render pathlines as tubes in 3d plots to minimize path point color artifacts
* use vtk with setup-headless-display-action instead of vtk-osmesa
* improve zoom and legend font size in 3d plots
* export 3d plots as PNG instead of PDF/SVG
* wrap long comments in P04
* try pinning nbconvert==7.16.3 to fix notebook thumbnails as suggested in spatialaudio/nbsphinx#796

.doc/requirements.rtd.txt

@@ -1,3 +1,4 @@
+nbconvert==7.16.3
 nbsphinx
 nbsphinx_link
 ipython

.github/workflows/ex-rtd.yml

@@ -56,6 +56,9 @@ jobs:
         with:
           python-version: 3.9
 
+      - name: Set up headless display
+        uses: pyvista/setup-headless-display-action@v2
+
       - name: Install Python packages
         working-directory: etc
         run: |
@@ -64,14 +67,6 @@ jobs:
           pip install -r requirements.usgs.txt
           python -m ipykernel install --name python_kernel --user
 
-      # PyVista doesn't like the default vtk
-      - name: OpenGL workaround
-        if: runner.os == 'Linux'
-        run: |
-          # referenced from https://github.com/pyvista/pyvista/blob/main/.github/workflows/vtk-pre-test.yml#L53
-          pip uninstall -y vtk
-          pip install --extra-index-url https://wheels.vtk.org trame vtk-osmesa
-
       - name: Update flopy MODFLOW 6 classes
         run: python -m flopy.mf6.utils.generate_classes --ref develop --no-backup
 

.github/workflows/ex-workflow.yml

@@ -129,6 +129,9 @@ jobs:
         with:
           python-version: ${{ matrix.python }}
 
+      - name: Set up headless display
+        uses: pyvista/setup-headless-display-action@v2
+
       - name: Install Python packages
         working-directory: etc
         run: |
@@ -137,14 +140,6 @@ jobs:
           pip install -r requirements.usgs.txt
           python -m ipykernel install --name python_kernel --user
 
-      # PyVista doesn't like the default vtk
-      - name: OpenGL workaround
-        if: runner.os == 'Linux'
-        run: |
-          # referenced from https://github.com/pyvista/pyvista/blob/main/.github/workflows/vtk-pre-test.yml#L53
-          pip uninstall -y vtk
-          pip install --extra-index-url https://wheels.vtk.org trame vtk-osmesa
-
       - name: Update flopy MODFLOW 6 classes
         run: python -m flopy.mf6.utils.generate_classes --ref develop --no-backup
 

doc/sections/ex-prt-mp7-p02.tex

@@ -27,7 +27,7 @@ \subsection{Example Results}
 	Three-dimensional perspective of pathlines and 1000-day points for subproblem 2A, and recharge points for subproblem 2B. Points are colored by layer.
 }
 	{fig:ex-prt-mp7-p02-paths-3d}
-	{../figures/mp7-p02-paths-3d.pdf}
+	{../figures/mp7-p02-paths-3d.png}
 \end{StandardFigure}
 
 \begin{StandardFigure}{

doc/sections/ex-prt-mp7-p03.tex

@@ -30,7 +30,7 @@ \subsection{Example Results}
 
 \begin{StandardFigure}{
     Three-dimensional perspective of pathlines and 2000-day points. Points are colored by destination.
-    }{fig:ex-prt-mp7-p03-paths-3d}{../figures/mp7-p03-paths-3d.pdf}
+    }{fig:ex-prt-mp7-p03-paths-3d}{../figures/mp7-p03-paths-3d.png}
 \end{StandardFigure}
 
 \begin{StandardFigure}{

etc/ci_create_examples_rst.py

@@ -191,9 +191,7 @@
 
             tag = ".. figure:: ../figures/"
             if tag in line:
-                line = line.replace(tag, ".. figure:: ../_images/").replace(
-                    ".pdf", ".svg"
-                )
+                line = line.replace(tag, ".. figure:: ../_images/")
 
             tag = ".. figure:: ../images/"
             if tag in line:
@@ -312,11 +310,7 @@
         file_name
         for file_name in os.listdir(src_dir)
         if os.path.isfile(os.path.join(src_dir, file_name))
-        and (
-            file_name.endswith(".png")
-            or file_name.endswith(".gif")
-            or file_name.endswith(".svg")
-        )
+        and (file_name.endswith(".png") or file_name.endswith(".gif"))
     ]
     for file_name in file_names:
         src = os.path.join(src_dir, file_name)

scripts/ex-gwe-prt.py

@@ -639,7 +639,7 @@ def plot_results(gwf_sim, gwe_sim, prt_sim, silent=True):
         axes[1].set_xlim(0, 2000)
         xticks = np.arange(0, 2100, 250)
         axes[1].set_xticks(xticks)
-        axes[1].set_ylabel("Temperature, $^{\circ}C$")
+        axes[1].set_ylabel(r"Temperature, $^{\circ}C$")
         axes[1].set_ylim(40, 80.0)
         yticks = np.arange(40, 81, 10)
         axes[1].set_yticks(yticks)

scripts/ex-prt-mp7-p01.py

@@ -476,7 +476,7 @@ def get_pathlines(mf6_path, mp7_path):
         .tt
     )
 
-    # determine which particles ended up in which capture area
+    # determine which particles ended up in which capture zone
     mf6pl["destzone"] = mf6pl[mf6pl.istatus > 1].izone
     mf6pl["dest"] = mf6pl.apply(
         lambda row: (
@@ -757,33 +757,45 @@ def plot_pathpoints_3d(gwf, pathlines, title):
     bed_mesh.rotate_y(-20, point=axes.origin, inplace=True)
     bed_mesh.rotate_x(20, point=axes.origin, inplace=True)
 
-    p = pv.Plotter(window_size=[500, 500])
-    p.add_title(title, font_size=5)
-    p.add_mesh(gwf_mesh, opacity=0.025, style="wireframe")
-    p.add_mesh(
-        prt_mesh,
-        scalars="k" if "k" in prt_mesh.point_data else "ilay",
-        cmap=["green", "gold", "red"],
-        point_size=2,
-        render_points_as_spheres=True,
-    )
-    p.add_mesh(riv_mesh, color="teal", opacity=0.2)
-    p.add_mesh(wel_mesh, color="red", opacity=0.3)
-    p.add_mesh(bed_mesh, color="tan", opacity=0.1)
-    p.remove_scalar_bar()
-    p.add_legend(
-        labels=[("Layer 1", "green"), ("Layer 2", "gold"), ("Layer 3", "red")],
-        bcolor="white",
-        face="r",
-        size=(0.1, 0.1),
-    )
+    def _plot(screenshot=False):
+        p = pv.Plotter(
+            window_size=[500, 500],
+            off_screen=screenshot,
+            notebook=False if screenshot else None,
+        )
+        p.enable_anti_aliasing()
+        p.add_title(title, font_size=7)
+        p.add_mesh(gwf_mesh, opacity=0.025, style="wireframe")
+        p.add_mesh(
+            prt_mesh,
+            scalars="k" if "k" in prt_mesh.point_data else "ilay",
+            cmap=["green", "gold", "red"],
+            point_size=4,
+            line_width=4,
+            render_points_as_spheres=True,
+            render_lines_as_tubes=True,
+            smooth_shading=True,
+        )
+        p.add_mesh(riv_mesh, color="teal", opacity=0.2)
+        p.add_mesh(wel_mesh, color="red", opacity=0.3)
+        p.add_mesh(bed_mesh, color="tan", opacity=0.1)
+        p.remove_scalar_bar()
+        p.add_legend(
+            labels=[("Layer 1", "green"), ("Layer 2", "gold"), ("Layer 3", "red")],
+            bcolor="white",
+            face="r",
+            size=(0.15, 0.15),
+        )
 
-    if plot_save:
-        p.save_graphic(figs_path / f"{sim_name}-paths-3d.pdf")
-        p.save_graphic(figs_path / f"{sim_name}-paths-3d.svg")
-    if plot_show:
         p.camera.zoom(1.7)
         p.show()
+        if screenshot:
+            p.screenshot(figs_path / f"{sim_name}-paths-3d.png")
+
+    if plot_show:
+        _plot()
+    if plot_save:
+        _plot(screenshot=True)
 
 
 def plot_all_pathlines(gwf, mf6pl, title=None):

scripts/ex-prt-mp7-p02.py

@@ -1085,47 +1085,64 @@ def plot_3d(gwf, pathlines, endpoints=None, title=None):
     bed_mesh.rotate_y(-10, point=axes.origin, inplace=True)
     bed_mesh.rotate_x(10, point=axes.origin, inplace=True)
 
-    p = pv.Plotter(window_size=[500, 500])
-    if title is not None:
-        p.add_title(title, font_size=5)
-    p.add_mesh(gwf_mesh, opacity=0.025, style="wireframe")
-    p.add_mesh(
-        prt_mesh,
-        scalars="k" if "k" in prt_mesh.point_data else "ilay",
-        cmap=["green", "gold", "red"],
-        point_size=2,
-    )
-    p.remove_scalar_bar()
+    plot_endpoints = False
     if endpoints is not None:
+        plot_endpoints = True
         endpoints = endpoints.to_records(index=False)
         endpoints["z"] = endpoints["z"] * vert_exag
         eps_mesh = pv.PolyData(np.array(tuple(map(tuple, endpoints[["x", "y", "z"]]))))
         eps_mesh.rotate_z(-30, point=axes.origin, inplace=True)
         eps_mesh.rotate_y(-10, point=axes.origin, inplace=True)
         eps_mesh.rotate_x(10, point=axes.origin, inplace=True)
+
+    def _plot(screenshot=False):
+        p = pv.Plotter(
+            window_size=[500, 500],
+            off_screen=screenshot,
+            notebook=False if screenshot else None,
+        )
+        p.enable_anti_aliasing()
+        if title is not None:
+            p.add_title(title, font_size=7)
+        p.add_mesh(gwf_mesh, opacity=0.025, style="wireframe")
         p.add_mesh(
-            eps_mesh,
-            scalars=endpoints.k.ravel(),
+            prt_mesh,
+            scalars="k" if "k" in prt_mesh.point_data else "ilay",
             cmap=["green", "gold", "red"],
-            point_size=2,
+            point_size=4,
+            line_width=3,
+            render_points_as_spheres=True,
+            render_lines_as_tubes=True,
+            smooth_shading=True,
         )
         p.remove_scalar_bar()
-    p.add_mesh(riv_mesh, color="teal", opacity=0.2)
-    p.add_mesh(wel_mesh, color="red", opacity=0.3)
-    p.add_mesh(bed_mesh, color="tan", opacity=0.1)
-    p.add_legend(
-        labels=[("Layer 1", "green"), ("Layer 2", "gold"), ("Layer 3", "red")],
-        bcolor="white",
-        face="r",
-        size=(0.1, 0.1),
-    )
+        if plot_endpoints:
+            p.add_mesh(
+                eps_mesh,
+                scalars=endpoints.k.ravel(),
+                cmap=["green", "gold", "red"],
+                point_size=4,
+            )
+            p.remove_scalar_bar()
+        p.add_mesh(riv_mesh, color="teal", opacity=0.2)
+        p.add_mesh(wel_mesh, color="red", opacity=0.3)
+        p.add_mesh(bed_mesh, color="tan", opacity=0.1)
+        p.add_legend(
+            labels=[("Layer 1", "green"), ("Layer 2", "gold"), ("Layer 3", "red")],
+            bcolor="white",
+            face="r",
+            size=(0.15, 0.15),
+        )
 
-    if plot_save:
-        p.save_graphic(figs_path / f"{sim_name}-paths-3d.pdf")
-        p.save_graphic(figs_path / f"{sim_name}-paths-3d.svg")
-    if plot_show:
-        p.camera.zoom(3.0)
+        p.camera.zoom(2)
         p.show()
+        if screenshot:
+            p.screenshot(figs_path / f"{sim_name}-paths-3d.png")
+
+    if plot_show:
+        _plot()
+    if plot_save:
+        _plot(screenshot=True)
 
 
 def load_head():

scripts/ex-prt-mp7-p03.py

@@ -524,7 +524,7 @@ def get_mf6_pathlines(path):
     # index temporarily by composite key fields
     pl.set_index(["iprp", "irpt", "trelease"], drop=False, inplace=True)
 
-    # determine which particles ended up in which capture area
+    # determine which particles ended up in which capture zone
     pl["destzone"] = pl[pl.istatus > 1].izone
     pl["dest"] = pl.apply(
         lambda row: (
@@ -541,12 +541,6 @@ def get_mf6_pathlines(path):
         axis=1,
     )
 
-    # add markercolor column, color-coding by layer for plots
-    pl["mc"] = pl.apply(
-        lambda row: "green" if row.ilay == 1 else "gold" if row.ilay == 2 else "red",
-        axis=1,
-    )
-
     # reset index
     pl.reset_index(drop=True, inplace=True)
 
@@ -785,40 +779,52 @@ def plot_pathpoints_3d(gwf, mf6pl, title=None):
     bed_mesh.rotate_y(-10, point=axes.origin, inplace=True)
     bed_mesh.rotate_x(10, point=axes.origin, inplace=True)
 
-    p = pv.Plotter(window_size=[500, 500])
-    if title is not None:
-        p.add_title(title, font_size=5)
-    p.add_mesh(gwf_mesh, opacity=0.025, style="wireframe")
-    p.add_mesh(
-        prt_mesh,
-        scalars="destzone",
-        cmap=["red", "red", "green", "blue"],
-        point_size=3,
-        render_points_as_spheres=True,
-    )
-    p.add_mesh(drn_mesh, color="green", opacity=0.2)
-    p.add_mesh(riv_mesh, color="teal", opacity=0.2)
-    p.add_mesh(wel_mesh, color="red", opacity=0.3)
-    p.add_mesh(wel2_mesh, color="red", opacity=0.2)
-    p.add_mesh(bed_mesh, color="tan", opacity=0.1)
-    p.remove_scalar_bar()
-    p.add_legend(
-        labels=[
-            ("Well (layer 3)", "red"),
-            ("Drain", "green"),
-            ("River", "blue"),
-        ],
-        bcolor="white",
-        face="r",
-        size=(0.1, 0.1),
-    )
+    def _plot(screenshot=False):
+        p = pv.Plotter(
+            window_size=[500, 500],
+            off_screen=screenshot,
+            notebook=False if screenshot else None,
+        )
+        p.enable_anti_aliasing()
+        if title is not None:
+            p.add_title(title, font_size=7)
+        p.add_mesh(gwf_mesh, opacity=0.025, style="wireframe")
+        p.add_mesh(
+            prt_mesh,
+            scalars="destzone",
+            cmap=["red", "red", "green", "blue"],
+            point_size=4,
+            line_width=3,
+            render_points_as_spheres=True,
+            render_lines_as_tubes=True,
+            smooth_shading=True,
+        )
+        p.add_mesh(drn_mesh, color="green", opacity=0.2)
+        p.add_mesh(riv_mesh, color="teal", opacity=0.2)
+        p.add_mesh(wel_mesh, color="red", opacity=0.3)
+        p.add_mesh(wel2_mesh, color="red", opacity=0.2)
+        p.add_mesh(bed_mesh, color="tan", opacity=0.1)
+        p.remove_scalar_bar()
+        p.add_legend(
+            labels=[
+                ("Well (layer 3)", "red"),
+                ("Drain", "green"),
+                ("River", "blue"),
+            ],
+            bcolor="white",
+            face="r",
+            size=(0.2, 0.2),
+        )
 
-    if plot_save:
-        p.save_graphic(figs_path / f"{sim_name}-paths-3d.pdf")
-        p.save_graphic(figs_path / f"{sim_name}-paths-3d.svg")
-    if plot_show:
-        p.camera.zoom(3.8)
+        p.camera.zoom(2.2)
         p.show()
+        if screenshot:
+            p.screenshot(figs_path / f"{sim_name}-paths-3d.png")
+
+    if plot_show:
+        _plot()
+    if plot_save:
+        _plot(screenshot=True)
 
 
 def plot_all_pathlines(gwf, mf6pl, title=None):
@@ -930,7 +936,7 @@ def plot_all(gwfsim):
     plot_pathpoints_3d(
         gwf,
         mf6pathlines,
-        title="Pathlines and 2000-day points,\ncolored by destination, 3D view",
+        title="Pathlines, 2000-day points,\ncolored by destination",
     )
     plot_endpoints(
         gwf,