Skip to content

Commit 6ad177c

Browse files
jacklovelljacklovell
authored
Merge pull request #81 from skuba31/80-plotting
80 plotting
2 parents 267d844 + 11a9fa5 commit 6ad177c

File tree

6 files changed

+269
-38
lines changed

6 files changed

+269
-38
lines changed

‎CHANGELOG.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ Release 1.3.0 (TBD)
77
* Support Cherab 1.5 (#74).
88
* Package demo files in distribution (#77)
99
* Switch build from setuptools to scikit-build-core (#79).
10+
* Added a submodule for plotting, deprecating `SOLPSMesh` methods (#80)
1011

1112
Release 1.2.1 (17 Feb 2023)
1213
-------------------

‎cherab/solps/mesh_geometry.py‎

Lines changed: 7 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -16,10 +16,9 @@
1616
#
1717
# See the Licence for the specific language governing permissions and limitations
1818
# under the Licence.
19+
from warnings import warn
1920

2021
import numpy as np
21-
import matplotlib.pyplot as plt
22-
from matplotlib.collections import PolyCollection
2322

2423

2524
class SOLPSMesh:
@@ -337,22 +336,9 @@ def plot_triangle_mesh(self, solps_data=None, ax=None):
337336
338337
:param solps_data: Data array defined on the SOLPS mesh
339338
"""
340-
if ax is None:
341-
_, ax = plt.subplots(constrained_layout=True)
342-
343-
verts = self.vertex_coordinates[self.triangles]
344-
if solps_data is None:
345-
collection_mesh = PolyCollection(verts, facecolor="none", edgecolor='b', linewidth=0.5)
346-
else:
347-
collection_mesh = PolyCollection(verts)
348-
collection_mesh.set_array(solps_data[self.triangle_to_grid_map[:, 0], self.triangle_to_grid_map[:, 1]])
349-
ax.add_collection(collection_mesh)
350-
ax.set_aspect(1)
351-
ax.set_xlim(self.mesh_extent["minr"], self.mesh_extent["maxr"])
352-
ax.set_ylim(self.mesh_extent["minz"], self.mesh_extent["maxz"])
353-
ax.set_xlabel("R [m]")
354-
ax.set_ylabel("Z [m]")
355-
339+
warn("plot_triangle_mesh method is deprecated, use functions from plot module.", DeprecationWarning)
340+
from .plot import plot_triangle_mesh
341+
ax = plot_triangle_mesh(self, solps_data, ax)
356342
return ax
357343

358344
def plot_quadrangle_mesh(self, solps_data=None, ax=None):
@@ -361,21 +347,7 @@ def plot_quadrangle_mesh(self, solps_data=None, ax=None):
361347
362348
:param solps_data: Data array defined on the SOLPS mesh
363349
"""
364-
365-
if ax is None:
366-
_, ax = plt.subplots(constrained_layout=True)
367-
368-
verts = self.vertex_coordinates[self.quadrangles]
369-
if solps_data is None:
370-
collection_mesh = PolyCollection(verts, facecolor="none", edgecolor='b', linewidth=0.5)
371-
else:
372-
collection_mesh = PolyCollection(verts)
373-
collection_mesh.set_array(solps_data[self.quadrangle_to_grid_map[:, 0], self.quadrangle_to_grid_map[:, 1]])
374-
ax.add_collection(collection_mesh)
375-
ax.set_aspect(1)
376-
ax.set_xlim(self.mesh_extent["minr"], self.mesh_extent["maxr"])
377-
ax.set_ylim(self.mesh_extent["minz"], self.mesh_extent["maxz"])
378-
ax.set_xlabel("R [m]")
379-
ax.set_ylabel("Z [m]")
380-
350+
warn("plot_quadrangle_mesh method is deprecated, use functions from plot module.", DeprecationWarning)
351+
from .plot import plot_quadrangle_mesh
352+
ax = plot_quadrangle_mesh(self, solps_data, ax)
381353
return ax

‎cherab/solps/plot.py‎

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
import matplotlib.pyplot as plt
2+
from matplotlib.collections import PolyCollection
3+
4+
5+
def plot_quadrangle_mesh(mesh, solps_data=None, ax=None):
6+
"""
7+
Plots the quadrangle mesh grid geometry to a matplotlib figure.
8+
9+
If solps_data is provided, it is used to colour the faces of the quadrangles in the mesh.
10+
If matplotlib axes are provided the collection is added to them them,
11+
otherwise a new figure and axes are created.
12+
13+
:param mesh: SOLPSMesh object
14+
:param solps_data: Data array defined on the SOLPS mesh (optional)
15+
:param ax: matplotlib axes (optional)
16+
:return: matplotlib axes
17+
"""
18+
if ax is None:
19+
_, ax = plt.subplots(constrained_layout=True)
20+
21+
if solps_data is None:
22+
collection_mesh = create_quadrangle_polycollection(mesh, facecolor="none", edgecolor='b', linewidth=0.5)
23+
else:
24+
collection_mesh = create_quadrangle_polycollection(mesh, solps_data)
25+
ax.add_collection(collection_mesh)
26+
27+
ax = _format_matplotlib_axes(ax, mesh)
28+
return ax
29+
30+
31+
def plot_triangle_mesh(mesh, solps_data=None, ax=None):
32+
"""
33+
Plots the triangle mesh grid geometry to a matplotlib figure.
34+
35+
If solps_data is provided, it is used to colour the faces of the triangles in the mesh.
36+
If matplotlib axes are provided the collection is added to them them,
37+
otherwise a new figure and axes are created.
38+
39+
:param mesh: SOLPSMesh object
40+
:param solps_data: Data array defined on the SOLPS mesh
41+
:param ax: matplotlib axes (optional)
42+
:return: matplotlib axes
43+
"""
44+
if ax is None:
45+
_, ax = plt.subplots(constrained_layout=True)
46+
47+
if solps_data is None:
48+
collection_mesh = create_triangle_polycollection(mesh, facecolor="none", edgecolor='b', linewidth=0.5)
49+
else:
50+
collection_mesh = create_triangle_polycollection(mesh, solps_data)
51+
ax.add_collection(collection_mesh)
52+
53+
ax = _format_matplotlib_axes(ax, mesh)
54+
return ax
55+
56+
57+
def create_quadrangle_polycollection(mesh, solps_data=None, **collection_kw):
58+
"""
59+
Creates a matplotlib PolyCollection object from the quadrangle mesh.
60+
61+
If solps_data is provided, it is used to colour the faces of the quadrangles in the mesh.
62+
63+
:param mesh: SOLPSMesh object
64+
:param solps_data: Data array defined on the SOLPS mesh
65+
:param collection_kw: Keyword arguments for the PolyCollection
66+
:return: matplotlib.collections.PolyCollection
67+
"""
68+
verts = mesh.vertex_coordinates[mesh.quadrangles]
69+
collection_mesh = PolyCollection(verts, **collection_kw)
70+
if solps_data is not None:
71+
collection_mesh.set_array(solps_data[mesh.quadrangle_to_grid_map[:, 0], mesh.quadrangle_to_grid_map[:, 1]])
72+
return collection_mesh
73+
74+
75+
def create_triangle_polycollection(mesh, solps_data=None, **collection_kw):
76+
"""
77+
Creates a matplotlib PolyCollection object from the triangle mesh.
78+
79+
If solps_data is provided, it is used to colour the faces of the triangles in the mesh.
80+
81+
:param mesh: SOLPSMesh object
82+
:param solps_data: Data array defined on the SOLPS mesh
83+
:param collection_kw: Keyword arguments for the PolyCollection
84+
:return: matplotlib.collections.PolyCollection
85+
"""
86+
verts = mesh.vertex_coordinates[mesh.triangles]
87+
collection_mesh = PolyCollection(verts, **collection_kw)
88+
if solps_data is not None:
89+
collection_mesh.set_array(solps_data[mesh.triangle_to_grid_map[:, 0], mesh.triangle_to_grid_map[:, 1]])
90+
return collection_mesh
91+
92+
93+
def _format_matplotlib_axes(ax, mesh=None):
94+
"""
95+
Formats the matplotlib axes for a SOLPS mesh plot.
96+
97+
Sets aspect and labels for the axes.
98+
If a SOLPSMesh object is provided, sets the limits of the axes to the mesh extent.
99+
100+
:param ax: matplotlib axes
101+
:param mesh: SOLPSMesh object (optional)
102+
:return: matplotlib axes
103+
"""
104+
ax.set_aspect(1)
105+
ax.set_xlabel("R [m]")
106+
ax.set_ylabel("z [m]")
107+
if mesh is not None:
108+
ax.set_xlim(mesh.mesh_extent["minr"], mesh.mesh_extent["maxr"])
109+
ax.set_ylim(mesh.mesh_extent["minz"], mesh.mesh_extent["maxz"])
110+
return ax

‎demos/plots/collections.py‎

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,89 @@
1+
# Copyright 2014-2020 United Kingdom Atomic Energy Authority
2+
#
3+
# Licensed under the EUPL, Version 1.1 or – as soon they will be approved by the
4+
# European Commission - subsequent versions of the EUPL (the "Licence");
5+
# You may not use this work except in compliance with the Licence.
6+
# You may obtain a copy of the Licence at:
7+
#
8+
# https://joinup.ec.europa.eu/software/page/eupl5
9+
#
10+
# Unless required by applicable law or agreed to in writing, software distributed
11+
# under the Licence is distributed on an "AS IS" basis, WITHOUT WARRANTIES OR
12+
# CONDITIONS OF ANY KIND, either express or implied.
13+
#
14+
# See the Licence for the specific language governing permissions and limitations
15+
# under the Licence.
16+
import os
17+
import matplotlib.pyplot as plt
18+
19+
from cherab.solps import load_solps_from_raw_output
20+
from cherab.solps.plot import create_quadrangle_polycollection, create_triangle_polycollection
21+
22+
23+
plt.rcParams['figure.figsize'] = [5, 10] # default figure size
24+
25+
# Load the simulation.
26+
demos_directory = os.path.dirname(os.path.dirname(__file__))
27+
simulation_directory = os.path.join(demos_directory, 'data', 'raw')
28+
print('Loading simulation...')
29+
sim = load_solps_from_raw_output(simulation_directory)
30+
mesh = sim.mesh
31+
32+
# plot quadrangle and triangle meshes
33+
# plot the quadrangle b2 mesh
34+
collection_qm = create_quadrangle_polycollection(mesh, facecolor="none", edgecolor='b', linewidth=0.2)
35+
36+
fig_qmesh, ax_qmesh = plt.subplots()
37+
ax_qmesh.set_title("Quadrangle B2 Mesh")
38+
ax_qmesh.add_collection(collection_qm)
39+
40+
ax_qmesh.set_aspect("equal")
41+
ax_qmesh.set_xlabel("R [m]")
42+
ax_qmesh.set_ylabel("z [m]")
43+
ax_qmesh.autoscale() # adding a collection does not change the limits of the axes
44+
45+
#plot the quadrangle b2 mesh with b2 ion temperature values
46+
collection_qti = create_quadrangle_polycollection(mesh, solps_data=sim.ion_temperature)
47+
48+
fig_qti, ax_qti = plt.subplots()
49+
ax_qti.set_title("B2 Ion Temperature")
50+
ax_qti.add_collection(collection_qti)
51+
cax_qti = ax_qti.inset_axes([1.05, 0, 0.05, 1])
52+
fig_qti.colorbar(collection_qti, cax=cax_qti, label="Ion Temperature [eV]")
53+
54+
ax_qti.set_aspect("equal")
55+
ax_qti.set_xlabel("R [m]")
56+
ax_qti.set_ylabel("z [m]")
57+
ax_qti.set_xlim(mesh.mesh_extent["minr"], mesh.mesh_extent["maxr"])
58+
ax_qti.set_ylim(mesh.mesh_extent["minz"], mesh.mesh_extent["maxz"])
59+
60+
# plot the triangle B2 mesh
61+
collection_tm = create_triangle_polycollection(mesh, facecolor="none", edgecolor='g', linewidth=0.25)
62+
63+
fig_tmesh, ax_tmesh = plt.subplots()
64+
ax_tmesh.set_title("Cherab Triangle Mesh")
65+
ax_tmesh.add_collection(collection_tm)
66+
67+
ax_tmesh.set_aspect("equal")
68+
ax_tmesh.set_xlabel("R [m]")
69+
ax_tmesh.set_ylabel("z [m]")
70+
ax_tmesh.set_xlim(mesh.mesh_extent["minr"], mesh.mesh_extent["maxr"])
71+
ax_tmesh.set_ylim(mesh.mesh_extent["minz"], mesh.mesh_extent["maxz"])
72+
73+
74+
# plot the triangle B2 mesh with b2 ion temperature values
75+
collection_tti = create_triangle_polycollection(mesh, solps_data=sim.ion_temperature, edgecolors='face')
76+
77+
fig_tti, ax_tti = plt.subplots()
78+
ax_tti.set_title("Cherab Triangle mesh with Ion Temperature")
79+
ax_tti.add_collection(collection_tti)
80+
cax_tti = ax_tti.inset_axes([1.05, 0, 0.05, 1])
81+
fig_tti.colorbar(collection_tti, cax=cax_tti, label="Ion Temperature [eV]")
82+
83+
ax_tti.set_aspect("equal")
84+
ax_tti.set_xlabel("R [m]")
85+
ax_tti.set_ylabel("z [m]")
86+
ax_tti.set_xlim(mesh.mesh_extent["minr"], mesh.mesh_extent["maxr"])
87+
ax_tti.set_ylim(mesh.mesh_extent["minz"], mesh.mesh_extent["maxz"])
88+
89+
plt.show()

‎demos/plots/mesh_and_values.py‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -20,30 +20,32 @@
2020
from matplotlib.collections import PolyCollection
2121

2222
from cherab.solps import load_solps_from_raw_output
23+
from cherab.solps.plot import plot_quadrangle_mesh, plot_triangle_mesh
2324

2425

2526
# Load the simulation.
2627
demos_directory = os.path.dirname(os.path.dirname(__file__))
2728
simulation_directory = os.path.join(demos_directory, 'data', 'raw')
2829
print('Loading simulation...')
2930
sim = load_solps_from_raw_output(simulation_directory)
31+
mesh = sim.mesh
3032

3133
# plot quadrangle and triangle meshes
3234
# plot the quadrangle b2 mesh
33-
ax = sim.mesh.plot_quadrangle_mesh()
35+
ax = plot_quadrangle_mesh(mesh)
3436
ax.set_title("Quadrangle B2 Mesh")
3537
ax.get_figure().set_size_inches((10, 20))
3638

3739
#plot the quadrangle b2 mesh with b2 ion temperature values
38-
ax = sim.mesh.plot_quadrangle_mesh(solps_data=sim.ion_temperature)
40+
ax = plot_quadrangle_mesh(mesh, solps_data=sim.ion_temperature)
3941
ax.get_figure().colorbar(ax.collections[0], aspect=40)
4042
ax.get_figure().set_size_inches((10, 20))
4143
ax.set_title("B2 Ion Temperature [eV]")
4244
plt.show()
4345

4446
# axes can also be passed as an argument
4547
fig_pass, ax = plt.subplots(figsize=(10, 20))
46-
ax = sim.mesh.plot_triangle_mesh(solps_data=sim.ion_temperature, ax=ax)
48+
ax = plot_triangle_mesh(mesh, solps_data=sim.ion_temperature, ax=ax)
4749
ax.get_figure().colorbar(ax.collections[0], aspect=40)
4850
ax.get_figure().set_size_inches((10, 20))
4951
ax.set_title("Cherab Triangle Mesh with Ion Temperature [eV]")

‎demos/plots/triplots.py‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
# Copyright 2014-2020 United Kingdom Atomic Energy Authority
2+
#
3+
# Licensed under the EUPL, Version 1.1 or – as soon they will be approved by the
4+
# European Commission - subsequent versions of the EUPL (the "Licence");
5+
# You may not use this work except in compliance with the Licence.
6+
# You may obtain a copy of the Licence at:
7+
#
8+
# https://joinup.ec.europa.eu/software/page/eupl5
9+
#
10+
# Unless required by applicable law or agreed to in writing, software distributed
11+
# under the Licence is distributed on an "AS IS" basis, WITHOUT WARRANTIES OR
12+
# CONDITIONS OF ANY KIND, either express or implied.
13+
#
14+
# See the Licence for the specific language governing permissions and limitations
15+
# under the Licence.
16+
import os
17+
18+
import matplotlib.pyplot as plt
19+
from matplotlib.tri import Triangulation
20+
21+
from cherab.solps import load_solps_from_raw_output
22+
23+
24+
# Load the simulation.
25+
demos_directory = os.path.dirname(os.path.dirname(__file__))
26+
simulation_directory = os.path.join(demos_directory, 'data', 'raw')
27+
print('Loading simulation...')
28+
29+
sim = load_solps_from_raw_output(simulation_directory)
30+
mesh = sim.mesh
31+
32+
# prepare data for triangulation plots using matplotlib.tri
33+
tri = Triangulation(mesh.vertex_coordinates[:, 0], mesh.vertex_coordinates[:, 1], mesh.triangles)
34+
ion_temperature_tri = sim.ion_temperature[mesh.triangle_to_grid_map[:, 0], mesh.triangle_to_grid_map[:, 1]]
35+
36+
37+
# plot mesh
38+
fig_mesh = plt.figure()
39+
ax_mesh = fig_mesh.add_subplot(111)
40+
ax_mesh.set_title("Mesh")
41+
ax_mesh.triplot(tri, lw=0.2)
42+
ax_mesh.set_aspect('equal')
43+
ax_mesh.set_xlabel("R [m]")
44+
ax_mesh.set_ylabel("z [m]")
45+
46+
47+
# plot ion temperature
48+
fig_ion_temperature = plt.figure()
49+
ax_ion_temperature = fig_ion_temperature.add_subplot(111)
50+
tpc = ax_ion_temperature.tripcolor(tri, ion_temperature_tri)
51+
fig_ion_temperature.colorbar(tpc, ax=ax_ion_temperature, label="Ion Temperature [eV]")
52+
ax_ion_temperature.set_title("Ion Temperature")
53+
ax_ion_temperature.set_aspect('equal')
54+
ax_ion_temperature.set_xlabel("R [m]")
55+
ax_ion_temperature.set_ylabel("z [m]")
56+
57+
plt.show()

0 commit comments

Comments
 (0)