80 plotting

cherab/solps/mesh_geometry.py

@@ -18,8 +18,6 @@
 # under the Licence.
 
 import numpy as np
-import matplotlib.pyplot as plt
-from matplotlib.collections import PolyCollection
 
 
 class SOLPSMesh:
@@ -337,22 +335,8 @@ def plot_triangle_mesh(self, solps_data=None, ax=None):
 
         :param solps_data: Data array defined on the SOLPS mesh
         """
-        if ax is None:
-            _, ax = plt.subplots(constrained_layout=True)
-
-        verts = self.vertex_coordinates[self.triangles]
-        if solps_data is None:
-            collection_mesh = PolyCollection(verts, facecolor="none", edgecolor='b', linewidth=0.5)
-        else:
-            collection_mesh = PolyCollection(verts)
-            collection_mesh.set_array(solps_data[self.triangle_to_grid_map[:, 0], self.triangle_to_grid_map[:, 1]])
-        ax.add_collection(collection_mesh)
-        ax.set_aspect(1)
-        ax.set_xlim(self.mesh_extent["minr"], self.mesh_extent["maxr"])
-        ax.set_ylim(self.mesh_extent["minz"], self.mesh_extent["maxz"])
-        ax.set_xlabel("R [m]")
-        ax.set_ylabel("Z [m]")
-
+        from .plotting import plot_triangle_mesh
+        ax = plot_triangle_mesh(self, solps_data, ax)
         return ax
 
     def plot_quadrangle_mesh(self, solps_data=None, ax=None):
@@ -361,21 +345,6 @@ def plot_quadrangle_mesh(self, solps_data=None, ax=None):
 
         :param solps_data: Data array defined on the SOLPS mesh
         """
-
-        if ax is None:
-            _, ax = plt.subplots(constrained_layout=True)
-
-        verts = self.vertex_coordinates[self.quadrangles]
-        if solps_data is None:
-            collection_mesh = PolyCollection(verts, facecolor="none", edgecolor='b', linewidth=0.5)
-        else:
-            collection_mesh = PolyCollection(verts)
-            collection_mesh.set_array(solps_data[self.quadrangle_to_grid_map[:, 0], self.quadrangle_to_grid_map[:, 1]])
-        ax.add_collection(collection_mesh)
-        ax.set_aspect(1)
-        ax.set_xlim(self.mesh_extent["minr"], self.mesh_extent["maxr"])
-        ax.set_ylim(self.mesh_extent["minz"], self.mesh_extent["maxz"])
-        ax.set_xlabel("R [m]")
-        ax.set_ylabel("Z [m]")
-
+        from .plotting import plot_quadrangle_mesh
+        ax = plot_quadrangle_mesh(self, solps_data, ax)
         return ax

cherab/solps/plotting.py

@@ -0,0 +1,100 @@
+import matplotlib.pyplot as plt
+import numpy as np
+from matplotlib.collections import PolyCollection
+
+
+def plot_quadrangle_mesh(mesh, solps_data=None, ax=None):
+    """
+    Plot the quadrangle mesh grid geometry to a matplotlib figure.
+
+    :param mesh: SOLPSMesh object
+    :param solps_data: Data array defined on the SOLPS mesh (optional)
+    :param ax: matplotlib axes (optional)
+    :return: matplotlib axes
+    """
+    if ax is None:
+        _, ax = plt.subplots(constrained_layout=True)
+
+    if solps_data is None:
+        collection_mesh = create_quadrangle_polycollection(mesh, facecolor="none", edgecolor='b', linewidth=0.5)
+    else:
+        collection_mesh = create_quadrangle_polycollection(mesh, solps_data)
+    ax.add_collection(collection_mesh)
+
+    ax = format_matplotlib_axes(ax, mesh)
+    return ax
+
+
+def create_quadrangle_polycollection(mesh, solps_data: np.ndarray = None, **collection_kw):
+    """
+    Creates a matplotlib PolyCollection object from the quadrangle mesh.
+
+    If solps_data is provided, it is used to colour the faces of the quadrangles in the mesh.
+
+    :param mesh: SOLPSMesh object
+    :param solps_data: Optional[np.ndarray] - Data array defined on the SOLPS mesh
+    :param collection_kw: Keyword arguments for the PolyCollection
+
+    :return: matplotlib.collections.PolyCollection
+    """
+    verts = mesh.vertex_coordinates[mesh.quadrangles]
+    collection_mesh = PolyCollection(verts, **collection_kw)
+    if solps_data is not None:
+        collection_mesh.set_array(solps_data[mesh.quadrangle_to_grid_map[:, 0], mesh.quadrangle_to_grid_map[:, 1]])
+    return collection_mesh
+
+
+def format_matplotlib_axes(ax: plt.Axes, mesh=None) -> plt.Axes:
+    """
+    Format the matplotlib axes for a SOLPS mesh plot.
+
+    :param ax: matplotlib axes
+    :param mesh: SOLPSMesh object (optional)
+    :return: matplotlib axes
+    """
+    ax.set_aspect(1)
+    ax.set_xlabel("R [m]")
+    ax.set_ylabel("z [m]")
+    if mesh is not None:
+        ax.set_xlim(mesh.mesh_extent["minr"], mesh.mesh_extent["maxr"])
+        ax.set_ylim(mesh.mesh_extent["minz"], mesh.mesh_extent["maxz"])
+    return ax
+
+
+def plot_triangle_mesh(mesh, solps_data: np.ndarray = None, ax: plt.Axes = None) -> plt.Axes:
+    """
+    Plot the triangle mesh grid geometry to a matplotlib figure.
+
+    :param mesh: SOLPSMesh object
+    :param solps_data: Data array defined on the SOLPS mesh
+    :param ax: matplotlib axes (optional)
+    :return: matplotlib axes
+    """
+    if ax is None:
+        _, ax = plt.subplots(constrained_layout=True)
+
+    if solps_data is None:
+        collection_mesh = create_triangle_polycollection(mesh, facecolor="none", edgecolor='b', linewidth=0.5)
+    else:
+        collection_mesh = create_triangle_polycollection(mesh, solps_data)
+    ax.add_collection(collection_mesh)
+
+    ax = format_matplotlib_axes(ax, mesh)
+    return ax
+
+
+def create_triangle_polycollection(mesh, solps_data: np.ndarray = None, **collection_kw) -> PolyCollection:
+    """
+    Creates a matplotlib PolyCollection object from the triangle mesh.
+
+    If solps_data is provided, it is used to colour the faces of the triangles in the mesh.
+
+    :param mesh: SOLPSMesh object
+    :param solps_data: Data array defined on the SOLPS mesh
+    :return: matplotlib.collections.PolyCollection
+    """
+    verts = mesh.vertex_coordinates[mesh.triangles]
+    collection_mesh = PolyCollection(verts, **collection_kw)
+    if solps_data is not None:
+        collection_mesh.set_array(solps_data[mesh.triangle_to_grid_map[:, 0], mesh.triangle_to_grid_map[:, 1]])
+    return collection_mesh

demos/plots/collections.py

@@ -0,0 +1,72 @@
+# Copyright 2014-2020 United Kingdom Atomic Energy Authority
+#
+# Licensed under the EUPL, Version 1.1 or – as soon they will be approved by the
+# European Commission - subsequent versions of the EUPL (the "Licence");
+# You may not use this work except in compliance with the Licence.
+# You may obtain a copy of the Licence at:
+#
+# https://joinup.ec.europa.eu/software/page/eupl5
+#
+# Unless required by applicable law or agreed to in writing, software distributed
+# under the Licence is distributed on an "AS IS" basis, WITHOUT WARRANTIES OR
+# CONDITIONS OF ANY KIND, either express or implied.
+#
+# See the Licence for the specific language governing permissions and limitations
+# under the Licence.
+import os
+import matplotlib.pyplot as plt
+
+from cherab.solps import load_solps_from_raw_output
+from cherab.solps.plotting import create_quadrangle_polycollection, create_triangle_polycollection, format_matplotlib_axes
+
+
+plt.rcParams['figure.figsize'] = [5, 10]  # default figure size
+
+# Load the simulation.
+demos_directory = os.path.dirname(os.path.dirname(__file__))
+simulation_directory = os.path.join(demos_directory, 'data', 'raw')
+print('Loading simulation...')
+sim = load_solps_from_raw_output(simulation_directory)
+mesh = sim.mesh
+
+# plot quadrangle and triangle meshes
+# plot the quadrangle b2 mesh
+collection_qm = create_quadrangle_polycollection(mesh, facecolor="none", edgecolor='b', linewidth=0.2)
+
+fig_qmesh = plt.figure()
+ax_qmesh = plt.subplot(111)
+ax_qmesh.set_title("Quadrangle B2 Mesh")
+ax_qmesh.add_collection(collection_qm)
+format_matplotlib_axes(ax_qmesh, mesh)
+
+#plot the quadrangle b2 mesh with b2 ion temperature values
+collection_qti = create_quadrangle_polycollection(mesh, solps_data=sim.ion_temperature)
+fig_qti = plt.figure()
+ax_qti = plt.subplot(111)
+ax_qti.set_title("B2 Ion Temperature")
+ax_qti.add_collection(collection_qti)
+cax_qti = ax_qti.inset_axes([1.05, 0, 0.05, 1])
+fig_qti.colorbar(collection_qti, cax=cax_qti, label="Ion Temperature [eV]")
+format_matplotlib_axes(ax_qti, mesh)
+
+# plot the triangle B2 mesh
+collection_tm = create_triangle_polycollection(mesh, facecolor="none", edgecolor='g', linewidth=0.25)
+
+fig_tmesh = plt.figure()
+ax_tmesh = plt.subplot(111)
+ax_tmesh.set_title("Cherab Triangle Mesh")
+ax_tmesh.add_collection(collection_tm)
+format_matplotlib_axes(ax_tmesh, mesh)
+
+# plot the triangle B2 mesh with b2 ion temperature values
+collection_tti = create_triangle_polycollection(mesh, solps_data=sim.ion_temperature, edgecolors='face')
+
+fig_tti = plt.figure()
+ax_tti = plt.subplot(111)
+ax_tti.set_title("Cherab Triangle mesh with Ion Temperature")
+ax_tti.add_collection(collection_tti)
+cax_tti = ax_tti.inset_axes([1.05, 0, 0.05, 1])
+fig_tti.colorbar(collection_tti, cax=cax_tti, label="Ion Temperature [eV]")
+format_matplotlib_axes(ax_tti, mesh)
+
+plt.show()

demos/plots/mesh_and_values.py

@@ -20,30 +20,32 @@
 from matplotlib.collections import PolyCollection
 
 from cherab.solps import load_solps_from_raw_output
+from cherab.solps.plotting import plot_quadrangle_mesh, plot_triangle_mesh
 
 
 # Load the simulation.
 demos_directory = os.path.dirname(os.path.dirname(__file__))
 simulation_directory = os.path.join(demos_directory, 'data', 'raw')
 print('Loading simulation...')
 sim = load_solps_from_raw_output(simulation_directory)
+mesh = sim.mesh
 
 # plot quadrangle and triangle meshes
 # plot the quadrangle b2 mesh
-ax = sim.mesh.plot_quadrangle_mesh()
+ax = plot_quadrangle_mesh(mesh)
 ax.set_title("Quadrangle B2 Mesh")
 ax.get_figure().set_size_inches((10, 20))
 
 #plot the quadrangle b2 mesh with b2 ion temperature values
-ax = sim.mesh.plot_quadrangle_mesh(solps_data=sim.ion_temperature)
+ax = plot_quadrangle_mesh(mesh, solps_data=sim.ion_temperature)
 ax.get_figure().colorbar(ax.collections[0], aspect=40)
 ax.get_figure().set_size_inches((10, 20))
 ax.set_title("B2 Ion Temperature [eV]")
 plt.show()
 
 # axes can also be passed as an argument
 fig_pass, ax = plt.subplots(figsize=(10, 20))
-ax = sim.mesh.plot_triangle_mesh(solps_data=sim.ion_temperature, ax=ax)
+ax = plot_triangle_mesh(mesh, solps_data=sim.ion_temperature, ax=ax)
 ax.get_figure().colorbar(ax.collections[0], aspect=40)
 ax.get_figure().set_size_inches((10, 20))
 ax.set_title("Cherab Triangle Mesh with Ion Temperature [eV]")

demos/plots/triplots.py

@@ -0,0 +1,53 @@
+# Copyright 2014-2020 United Kingdom Atomic Energy Authority
+#
+# Licensed under the EUPL, Version 1.1 or – as soon they will be approved by the
+# European Commission - subsequent versions of the EUPL (the "Licence");
+# You may not use this work except in compliance with the Licence.
+# You may obtain a copy of the Licence at:
+#
+# https://joinup.ec.europa.eu/software/page/eupl5
+#
+# Unless required by applicable law or agreed to in writing, software distributed
+# under the Licence is distributed on an "AS IS" basis, WITHOUT WARRANTIES OR
+# CONDITIONS OF ANY KIND, either express or implied.
+#
+# See the Licence for the specific language governing permissions and limitations
+# under the Licence.
+import os
+
+import matplotlib.pyplot as plt
+from matplotlib.tri import Triangulation
+
+from cherab.solps import load_solps_from_raw_output, SOLPSSimulation
+from cherab.solps.plotting import format_matplotlib_axes
+
+
+# Load the simulation.
+demos_directory = os.path.dirname(os.path.dirname(__file__))
+simulation_directory = os.path.join(demos_directory, 'data', 'raw')
+print('Loading simulation...')
+
+sim: SOLPSSimulation = load_solps_from_raw_output(simulation_directory)
+mesh = sim.mesh
+
+# prepare data for triangulation plots using matplotlib.tri
+tri = Triangulation(mesh.vertex_coordinates[:, 0], mesh.vertex_coordinates[:, 1], mesh.triangles)
+ion_temperature_tri = sim.ion_temperature[mesh.triangle_to_grid_map[:, 0], mesh.triangle_to_grid_map[:, 1]]
+
+
+# plot mesh
+fig_mesh = plt.figure()
+ax_mesh = fig_mesh.add_subplot(111)
+ax_mesh.set_title("Mesh")
+ax_mesh.triplot(tri, lw=0.2)
+format_matplotlib_axes(ax_mesh, mesh)
+
+
+# plot ion temperature
+fig_ion_temperature = plt.figure()
+ax_ion_temperature = fig_ion_temperature.add_subplot(111)
+tpc = ax_ion_temperature.tripcolor(tri, ion_temperature_tri)
+fig_ion_temperature.colorbar(tpc, ax=ax_ion_temperature, label="Ion Temperature [eV]")
+format_matplotlib_axes(ax_ion_temperature, mesh)
+
+plt.show()