From 330a86c9cd7b2bac8c86464c8e46fa00b0d629d1 Mon Sep 17 00:00:00 2001 From: svoboda Date: Tue, 14 Jan 2025 13:53:51 +0100 Subject: [PATCH 1/9] Add plotting submodule, remove plotting from mesh class --- cherab/solps/mesh_geometry.py | 39 ++----------- cherab/solps/plotting.py | 100 +++++++++++++++++++++++++++++++++ demos/plots/mesh_and_values.py | 8 ++- 3 files changed, 109 insertions(+), 38 deletions(-) create mode 100644 cherab/solps/plotting.py diff --git a/cherab/solps/mesh_geometry.py b/cherab/solps/mesh_geometry.py index 6704f71..9fd279a 100755 --- a/cherab/solps/mesh_geometry.py +++ b/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 diff --git a/cherab/solps/plotting.py b/cherab/solps/plotting.py new file mode 100644 index 0000000..64edfda --- /dev/null +++ b/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 diff --git a/demos/plots/mesh_and_values.py b/demos/plots/mesh_and_values.py index 159ea5d..1efa91b 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.plotting 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]") From 90dbb5edff7e795ef68358a9300b1ccf5f6144b8 Mon Sep 17 00:00:00 2001 From: svoboda Date: Tue, 14 Jan 2025 13:54:17 +0100 Subject: [PATCH 2/9] Add triplot and collections tutorials --- demos/plots/collections.py | 72 ++++++++++++++++++++++++++++++++++++++ demos/plots/triplots.py | 53 ++++++++++++++++++++++++++++ 2 files changed, 125 insertions(+) create mode 100644 demos/plots/collections.py create mode 100644 demos/plots/triplots.py diff --git a/demos/plots/collections.py b/demos/plots/collections.py new file mode 100644 index 0000000..7b3f6bf --- /dev/null +++ b/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() diff --git a/demos/plots/triplots.py b/demos/plots/triplots.py new file mode 100644 index 0000000..580f59f --- /dev/null +++ b/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() From 3387e9c9aaaea55b2c0b8c6863136601c254153e Mon Sep 17 00:00:00 2001 From: svoboda Date: Tue, 14 Jan 2025 18:08:50 +0100 Subject: [PATCH 3/9] Rename to plot, add deprecation warnings, remove typing --- CHANGELOG.md | 4 ++++ cherab/solps/mesh_geometry.py | 7 +++++-- cherab/solps/{plotting.py => plot.py} | 26 ++++++++++++++++++-------- demos/plots/collections.py | 2 +- demos/plots/mesh_and_values.py | 2 +- demos/plots/triplots.py | 6 +++--- 6 files changed, 32 insertions(+), 15 deletions(-) rename cherab/solps/{plotting.py => plot.py} (75%) diff --git a/CHANGELOG.md b/CHANGELOG.md index ecb7b33..de88131 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,10 @@ Project Changelog ================= +Release 1.2.2 +------------- +* Added a submodule for plotting, deprecating SOLPSMesh methods + Release 1.2.1 (17 Feb 2023) ------------------- diff --git a/cherab/solps/mesh_geometry.py b/cherab/solps/mesh_geometry.py index 9fd279a..2f61980 100755 --- a/cherab/solps/mesh_geometry.py +++ b/cherab/solps/mesh_geometry.py @@ -16,6 +16,7 @@ # # See the Licence for the specific language governing permissions and limitations # under the Licence. +from warnings import warn import numpy as np @@ -335,7 +336,8 @@ def plot_triangle_mesh(self, solps_data=None, ax=None): :param solps_data: Data array defined on the SOLPS mesh """ - from .plotting import plot_triangle_mesh + 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 @@ -345,6 +347,7 @@ def plot_quadrangle_mesh(self, solps_data=None, ax=None): :param solps_data: Data array defined on the SOLPS mesh """ - from .plotting import plot_quadrangle_mesh + 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/plotting.py b/cherab/solps/plot.py similarity index 75% rename from cherab/solps/plotting.py rename to cherab/solps/plot.py index 64edfda..01b344e 100644 --- a/cherab/solps/plotting.py +++ b/cherab/solps/plot.py @@ -1,11 +1,14 @@ 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. + 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) @@ -25,7 +28,7 @@ def plot_quadrangle_mesh(mesh, solps_data=None, ax=None): return ax -def create_quadrangle_polycollection(mesh, solps_data: np.ndarray = None, **collection_kw): +def create_quadrangle_polycollection(mesh, solps_data=None, **collection_kw): """ Creates a matplotlib PolyCollection object from the quadrangle mesh. @@ -44,9 +47,12 @@ def create_quadrangle_polycollection(mesh, solps_data: np.ndarray = None, **coll return collection_mesh -def format_matplotlib_axes(ax: plt.Axes, mesh=None) -> plt.Axes: +def format_matplotlib_axes(ax, mesh=None): """ - Format the matplotlib axes for a SOLPS mesh plot. + 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) @@ -61,9 +67,13 @@ def format_matplotlib_axes(ax: plt.Axes, mesh=None) -> plt.Axes: return ax -def plot_triangle_mesh(mesh, solps_data: np.ndarray = None, ax: plt.Axes = None) -> plt.Axes: +def plot_triangle_mesh(mesh, solps_data=None, ax=None): """ - Plot the triangle mesh grid geometry to a matplotlib figure. + 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 @@ -83,7 +93,7 @@ def plot_triangle_mesh(mesh, solps_data: np.ndarray = None, ax: plt.Axes = None) return ax -def create_triangle_polycollection(mesh, solps_data: np.ndarray = None, **collection_kw) -> PolyCollection: +def create_triangle_polycollection(mesh, solps_data=None, **collection_kw): """ Creates a matplotlib PolyCollection object from the triangle mesh. diff --git a/demos/plots/collections.py b/demos/plots/collections.py index 7b3f6bf..53cfa83 100644 --- a/demos/plots/collections.py +++ b/demos/plots/collections.py @@ -17,7 +17,7 @@ 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 +from cherab.solps.plot import create_quadrangle_polycollection, create_triangle_polycollection, format_matplotlib_axes plt.rcParams['figure.figsize'] = [5, 10] # default figure size diff --git a/demos/plots/mesh_and_values.py b/demos/plots/mesh_and_values.py index 1efa91b..b1b62b9 100644 --- a/demos/plots/mesh_and_values.py +++ b/demos/plots/mesh_and_values.py @@ -20,7 +20,7 @@ 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 +from cherab.solps.plot import plot_quadrangle_mesh, plot_triangle_mesh # Load the simulation. diff --git a/demos/plots/triplots.py b/demos/plots/triplots.py index 580f59f..4679c0a 100644 --- a/demos/plots/triplots.py +++ b/demos/plots/triplots.py @@ -18,8 +18,8 @@ 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 +from cherab.solps import load_solps_from_raw_output +from cherab.solps.plot import format_matplotlib_axes # Load the simulation. @@ -27,7 +27,7 @@ simulation_directory = os.path.join(demos_directory, 'data', 'raw') print('Loading simulation...') -sim: SOLPSSimulation = load_solps_from_raw_output(simulation_directory) +sim = load_solps_from_raw_output(simulation_directory) mesh = sim.mesh # prepare data for triangulation plots using matplotlib.tri From df23efc7caa10579e5737a3e5d573208798d1b6f Mon Sep 17 00:00:00 2001 From: svoboda Date: Wed, 15 Jan 2025 18:58:27 +0100 Subject: [PATCH 4/9] Add issue reference to CHANGELOG --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index de88131..477ecdd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,7 +3,7 @@ Project Changelog Release 1.2.2 ------------- -* Added a submodule for plotting, deprecating SOLPSMesh methods +* Added a submodule for plotting, deprecating `SOLPSMesh` methods ([#80](https://github.com/cherab/solps/issues/80)) Release 1.2.1 (17 Feb 2023) ------------------- From 9169fc38080a209cb3c7c8435e578736e02a4e43 Mon Sep 17 00:00:00 2001 From: svoboda Date: Thu, 16 Jan 2025 14:04:12 +0100 Subject: [PATCH 5/9] Remove function for formatting axes from tutorials. --- demos/plots/collections.py | 43 ++++++++++++++++++++++++++------------ demos/plots/triplots.py | 10 ++++++--- 2 files changed, 37 insertions(+), 16 deletions(-) diff --git a/demos/plots/collections.py b/demos/plots/collections.py index 53cfa83..8c35909 100644 --- a/demos/plots/collections.py +++ b/demos/plots/collections.py @@ -17,7 +17,7 @@ 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, format_matplotlib_axes +from cherab.solps.plot import create_quadrangle_polycollection, create_triangle_polycollection plt.rcParams['figure.figsize'] = [5, 10] # default figure size @@ -33,40 +33,57 @@ # 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) +fig_qmesh, ax_qmesh = plt.subplots() ax_qmesh.set_title("Quadrangle B2 Mesh") ax_qmesh.add_collection(collection_qm) -format_matplotlib_axes(ax_qmesh, mesh) + +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 = plt.figure() -ax_qti = plt.subplot(111) + +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]") -format_matplotlib_axes(ax_qti, mesh) + +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 = plt.figure() -ax_tmesh = plt.subplot(111) +fig_tmesh, ax_tmesh = plt.subplots() ax_tmesh.set_title("Cherab Triangle Mesh") ax_tmesh.add_collection(collection_tm) -format_matplotlib_axes(ax_tmesh, mesh) + +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 = plt.figure() -ax_tti = plt.subplot(111) +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]") -format_matplotlib_axes(ax_tti, mesh) + +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/triplots.py b/demos/plots/triplots.py index 4679c0a..b66c5fb 100644 --- a/demos/plots/triplots.py +++ b/demos/plots/triplots.py @@ -19,7 +19,6 @@ from matplotlib.tri import Triangulation from cherab.solps import load_solps_from_raw_output -from cherab.solps.plot import format_matplotlib_axes # Load the simulation. @@ -40,7 +39,9 @@ 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) +ax_mesh.set_aspect('equal') +ax_mesh.set_xlabel("R [m]") +ax_mesh.set_ylabel("z [m]") # plot ion temperature @@ -48,6 +49,9 @@ 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) +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() From 9bd9c9ecd370f8dab054461b39553a1ede7068cd Mon Sep 17 00:00:00 2001 From: svoboda Date: Thu, 16 Jan 2025 14:05:07 +0100 Subject: [PATCH 6/9] Make axes formatting function private to avoid micromanaging users' plots --- cherab/solps/plot.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/cherab/solps/plot.py b/cherab/solps/plot.py index 01b344e..f5e2b57 100644 --- a/cherab/solps/plot.py +++ b/cherab/solps/plot.py @@ -24,7 +24,7 @@ def plot_quadrangle_mesh(mesh, solps_data=None, ax=None): collection_mesh = create_quadrangle_polycollection(mesh, solps_data) ax.add_collection(collection_mesh) - ax = format_matplotlib_axes(ax, mesh) + ax = _format_matplotlib_axes(ax, mesh) return ax @@ -47,7 +47,7 @@ def create_quadrangle_polycollection(mesh, solps_data=None, **collection_kw): return collection_mesh -def format_matplotlib_axes(ax, mesh=None): +def _format_matplotlib_axes(ax, mesh=None): """ Formats the matplotlib axes for a SOLPS mesh plot. @@ -89,7 +89,7 @@ def plot_triangle_mesh(mesh, solps_data=None, ax=None): collection_mesh = create_triangle_polycollection(mesh, solps_data) ax.add_collection(collection_mesh) - ax = format_matplotlib_axes(ax, mesh) + ax = _format_matplotlib_axes(ax, mesh) return ax From 1f11e1e2c341adc7fdabf5f0d887e7bd505b9794 Mon Sep 17 00:00:00 2001 From: svoboda Date: Thu, 16 Jan 2025 14:11:11 +0100 Subject: [PATCH 7/9] Reorganize functions so that high level one are at the top --- cherab/solps/plot.py | 78 ++++++++++++++++++++++---------------------- 1 file changed, 39 insertions(+), 39 deletions(-) diff --git a/cherab/solps/plot.py b/cherab/solps/plot.py index f5e2b57..098a228 100644 --- a/cherab/solps/plot.py +++ b/cherab/solps/plot.py @@ -28,45 +28,6 @@ def plot_quadrangle_mesh(mesh, solps_data=None, ax=None): 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: 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, 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 - - def plot_triangle_mesh(mesh, solps_data=None, ax=None): """ Plots the triangle mesh grid geometry to a matplotlib figure. @@ -93,6 +54,25 @@ def plot_triangle_mesh(mesh, solps_data=None, ax=None): 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: 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 create_triangle_polycollection(mesh, solps_data=None, **collection_kw): """ Creates a matplotlib PolyCollection object from the triangle mesh. @@ -108,3 +88,23 @@ def create_triangle_polycollection(mesh, solps_data=None, **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 From 9af6f81c4b059600c6fe04c5a87cee17bdf3b280 Mon Sep 17 00:00:00 2001 From: svoboda Date: Thu, 16 Jan 2025 14:11:29 +0100 Subject: [PATCH 8/9] Fix version number in the changelog --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 477ecdd..7b23552 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,7 +1,7 @@ Project Changelog ================= -Release 1.2.2 +Release 1.3.0 ------------- * Added a submodule for plotting, deprecating `SOLPSMesh` methods ([#80](https://github.com/cherab/solps/issues/80)) From 10b1203fedfa6ce85bb3e5dce3d8cb97ac5f8161 Mon Sep 17 00:00:00 2001 From: svoboda Date: Mon, 27 Jan 2025 13:45:12 +0100 Subject: [PATCH 9/9] Fix and unify docstrings --- cherab/solps/plot.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cherab/solps/plot.py b/cherab/solps/plot.py index 098a228..215e82b 100644 --- a/cherab/solps/plot.py +++ b/cherab/solps/plot.py @@ -61,9 +61,8 @@ def create_quadrangle_polycollection(mesh, solps_data=None, **collection_kw): 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 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] @@ -81,6 +80,7 @@ def create_triangle_polycollection(mesh, solps_data=None, **collection_kw): :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]