diff --git a/CHANGELOG.md b/CHANGELOG.md index 4a104bc..8db2128 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,7 @@ Release 1.3.0 (TBD) * Support Cherab 1.5 (#74). * Package demo files in distribution (#77) * Switch build from setuptools to scikit-build-core (#79). +* Added a submodule for plotting, deprecating `SOLPSMesh` methods (#80) Release 1.2.1 (17 Feb 2023) ------------------- diff --git a/cherab/solps/mesh_geometry.py b/cherab/solps/mesh_geometry.py index 6704f71..2f61980 100755 --- a/cherab/solps/mesh_geometry.py +++ b/cherab/solps/mesh_geometry.py @@ -16,10 +16,9 @@ # # See the Licence for the specific language governing permissions and limitations # under the Licence. +from warnings import warn import numpy as np -import matplotlib.pyplot as plt -from matplotlib.collections import PolyCollection class SOLPSMesh: @@ -337,22 +336,9 @@ 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]") - + warn("plot_triangle_mesh method is deprecated, use functions from plot module.", DeprecationWarning) + from .plot 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 +347,7 @@ 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]") - + warn("plot_quadrangle_mesh method is deprecated, use functions from plot module.", DeprecationWarning) + from .plot import plot_quadrangle_mesh + ax = plot_quadrangle_mesh(self, solps_data, ax) return ax diff --git a/cherab/solps/plot.py b/cherab/solps/plot.py new file mode 100644 index 0000000..215e82b --- /dev/null +++ b/cherab/solps/plot.py @@ -0,0 +1,110 @@ +import matplotlib.pyplot as plt +from matplotlib.collections import PolyCollection + + +def plot_quadrangle_mesh(mesh, solps_data=None, ax=None): + """ + Plots the quadrangle mesh grid geometry to a matplotlib figure. + + If solps_data is provided, it is used to colour the faces of the quadrangles in the mesh. + If matplotlib axes are provided the collection is added to them them, + otherwise a new figure and axes are created. + + :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 plot_triangle_mesh(mesh, solps_data=None, ax=None): + """ + Plots the triangle mesh grid geometry to a matplotlib figure. + + If solps_data is provided, it is used to colour the faces of the triangles in the mesh. + If matplotlib axes are provided the collection is added to them them, + otherwise a new figure and axes are created. + + :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_quadrangle_polycollection(mesh, solps_data=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: 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 create_triangle_polycollection(mesh, solps_data=None, **collection_kw): + """ + 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 + :param collection_kw: Keyword arguments for the PolyCollection + :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 + + +def _format_matplotlib_axes(ax, mesh=None): + """ + Formats the matplotlib axes for a SOLPS mesh plot. + + Sets aspect and labels for the axes. + If a SOLPSMesh object is provided, sets the limits of the axes to the mesh extent. + + :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 diff --git a/demos/plots/collections.py b/demos/plots/collections.py new file mode 100644 index 0000000..8c35909 --- /dev/null +++ b/demos/plots/collections.py @@ -0,0 +1,89 @@ +# 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.plot import create_quadrangle_polycollection, create_triangle_polycollection + + +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, ax_qmesh = plt.subplots() +ax_qmesh.set_title("Quadrangle B2 Mesh") +ax_qmesh.add_collection(collection_qm) + +ax_qmesh.set_aspect("equal") +ax_qmesh.set_xlabel("R [m]") +ax_qmesh.set_ylabel("z [m]") +ax_qmesh.autoscale() # adding a collection does not change the limits of the axes + +#plot the quadrangle b2 mesh with b2 ion temperature values +collection_qti = create_quadrangle_polycollection(mesh, solps_data=sim.ion_temperature) + +fig_qti, ax_qti = plt.subplots() +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]") + +ax_qti.set_aspect("equal") +ax_qti.set_xlabel("R [m]") +ax_qti.set_ylabel("z [m]") +ax_qti.set_xlim(mesh.mesh_extent["minr"], mesh.mesh_extent["maxr"]) +ax_qti.set_ylim(mesh.mesh_extent["minz"], mesh.mesh_extent["maxz"]) + +# plot the triangle B2 mesh +collection_tm = create_triangle_polycollection(mesh, facecolor="none", edgecolor='g', linewidth=0.25) + +fig_tmesh, ax_tmesh = plt.subplots() +ax_tmesh.set_title("Cherab Triangle Mesh") +ax_tmesh.add_collection(collection_tm) + +ax_tmesh.set_aspect("equal") +ax_tmesh.set_xlabel("R [m]") +ax_tmesh.set_ylabel("z [m]") +ax_tmesh.set_xlim(mesh.mesh_extent["minr"], mesh.mesh_extent["maxr"]) +ax_tmesh.set_ylim(mesh.mesh_extent["minz"], mesh.mesh_extent["maxz"]) + + +# 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, ax_tti = plt.subplots() +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]") + +ax_tti.set_aspect("equal") +ax_tti.set_xlabel("R [m]") +ax_tti.set_ylabel("z [m]") +ax_tti.set_xlim(mesh.mesh_extent["minr"], mesh.mesh_extent["maxr"]) +ax_tti.set_ylim(mesh.mesh_extent["minz"], mesh.mesh_extent["maxz"]) + +plt.show() diff --git a/demos/plots/mesh_and_values.py b/demos/plots/mesh_and_values.py index 159ea5d..b1b62b9 100644 --- a/demos/plots/mesh_and_values.py +++ b/demos/plots/mesh_and_values.py @@ -20,6 +20,7 @@ from matplotlib.collections import PolyCollection from cherab.solps import load_solps_from_raw_output +from cherab.solps.plot import plot_quadrangle_mesh, plot_triangle_mesh # Load the simulation. @@ -27,15 +28,16 @@ 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]") @@ -43,7 +45,7 @@ # 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]") diff --git a/demos/plots/triplots.py b/demos/plots/triplots.py new file mode 100644 index 0000000..b66c5fb --- /dev/null +++ b/demos/plots/triplots.py @@ -0,0 +1,57 @@ +# 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 + + +# 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 + +# 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) +ax_mesh.set_aspect('equal') +ax_mesh.set_xlabel("R [m]") +ax_mesh.set_ylabel("z [m]") + + +# 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]") +ax_ion_temperature.set_title("Ion Temperature") +ax_ion_temperature.set_aspect('equal') +ax_ion_temperature.set_xlabel("R [m]") +ax_ion_temperature.set_ylabel("z [m]") + +plt.show()