Module traceon.plotting

The traceon.plotting module uses the vedo plotting library to provide some convenience functions to show the line and triangle meshes generated by Traceon.

To show a mesh, for example use:

plt.new_figure()
plt.plot_mesh(mesh)
plt.show()

Where mesh is created using the traceon.geometry module.

Functions

def get_current_figure() ‑> Figure

Expand source code

def get_current_figure() -> Figure:
    """Get the currently active figure. If no figure has been created yet
    a new figure will be returned.

    Returns
    --------------------
    `Figure`"""
    if len(_current_figures):
        return _current_figures[-1]
    
    return new_figure()

Get the currently active figure. If no figure has been created yet a new figure will be returned.

Returns

Figure

def new_figure(show_legend=True)

Expand source code

def new_figure(show_legend=True):
    """Create a new figure and make it the current figure.
    
    Parameters
    ----------------------
    show_legend: bool
        Whether to show the legend in the corner of the figure

    Returns
    ----------------------
    `Figure`"""
    global _current_figures
    f = Figure(show_legend=show_legend)
    _current_figures.append(f)
    return f

Create a new figure and make it the current figure.

Parameters

show_legend : bool

Whether to show the legend in the corner of the figure

Returns

Figure

def plot_charge_density(*args, **kwargs)

Expand source code

def plot_charge_density(*args, **kwargs):
    """Calls `Figure.plot_charge_density` on the current `Figure`"""
    get_current_figure().plot_charge_density(*args, **kwargs)

Calls Figure.plot_charge_density() on the current Figure

def plot_equipotential_lines(*args, **kwargs)

Expand source code

def plot_equipotential_lines(*args, **kwargs):
    """Calls `Figure.plot_equipotential_lines` on the current `Figure`"""
    get_current_figure().plot_equipotential_lines(*args, **kwargs)

Calls Figure.plot_equipotential_lines() on the current Figure

def plot_mesh(*args, **kwargs)

Expand source code

def plot_mesh(*args, **kwargs):
    """Calls `Figure.plot_mesh` on the current `Figure`"""
    get_current_figure().plot_mesh(*args, **kwargs)

Calls Figure.plot_mesh() on the current Figure

def plot_trajectories(*args, **kwargs)

Expand source code

def plot_trajectories(*args, **kwargs):
    """Calls `Figure.plot_trajectories` on the current `Figure`"""
    get_current_figure().plot_trajectories(*args, **kwargs)

Calls Figure.plot_trajectories() on the current Figure

def show()

Expand source code

def show():
    """Calls `Figure.show` on the current `Figure`"""
    global _current_figures
        
    for f in _current_figures:
        f.show()

    _current_figures = []

Calls Figure.show() on the current Figure

Classes

class Figure (show_legend=True)

Expand source code

class Figure:
    def __init__(self, show_legend=True):
        self.show_legend = show_legend
        self.is_2d = True
        self.legend_entries = []
        self.to_plot = []
     
    def plot_mesh(self, mesh, show_normals=False, **colors):
        """Plot mesh using the Vedo library. Optionally showing normal vectors.

        Parameters
        ---------------------
        mesh: `traceon.mesher.Mesh`
            The mesh to plot
        show_normals: bool
            Whether to show the normal vectors at every element
        colors: dict of (string, string)
            Use keyword arguments to specify colors, for example `plot_mesh(mesh, lens='blue', ground='green')`
        """
        if not len(mesh.triangles) and not len(mesh.lines):
            raise RuntimeError("Trying to plot empty mesh.")

        triangle_normals, line_normals = None, None
        
        if len(mesh.triangles):
            meshes, triangle_normals = _get_vedo_triangles_and_normals(mesh, **colors)
            self.legend_entries.extend(meshes)
            self.to_plot.append(meshes)
        
        if len(mesh.lines):
            lines, line_normals = _get_vedo_lines_and_normals(mesh, **colors)
            self.legend_entries.extend(lines)
            self.to_plot.append(lines)
         
        if show_normals:
            if triangle_normals is not None:
                self.to_plot.append(triangle_normals)
            if line_normals is not None:
                self.to_plot.append(line_normals)
        
        self.is_2d &= mesh.is_2d()

    def plot_equipotential_lines(self, field, surface, N0=75, N1=75, color_map='coolwarm', N_isolines=40, isolines_width=1, isolines_color='#444444'):
        """Make potential color plot including equipotential lines.

        Parameters
        -------------------------------------
        field: `traceon.solver.Field`
            The field used to compute the potential values (note that any field returned from the solver can be used)
        surface: `traceon.geometry.Surface`
            The surface in 3D space which will be 'colored in'
        N0: int
            Number of pixels to use along the first 'axis' of the surface
        N1: int
            Number of pixels to use along the second 'axis' of the surface
        color_map: str
            Color map to use to color in the surface
        N_isolines: int
            Number of equipotential lines to plot
        isolines_width: int
            The width to use for the isolines. Pass in 0 to disable the isolines.
        isolines_color: str
            Color to use for the isolines"""
        grid = _get_vedo_grid(field, surface, N0, N1)
        isolines = grid.isolines(n=N_isolines).color(isolines_color).lw(isolines_width) # type: ignore
        grid.cmap(color_map)
        self.to_plot.append(grid)
        self.to_plot.append(isolines)
    
    def plot_trajectories(self, trajectories, 
                xmin=None, xmax=None,
                ymin=None, ymax=None,
                zmin=None, zmax=None,
                color='#00AA00', line_width=1):
        """Plot particle trajectories.

        Parameters
        ------------------------------------
        trajectories: list of numpy.ndarray
            List of positions as returned by `traceon.tracing.Tracer.__call__`
        xmin, xmax: float
            Only plot trajectory points for which xmin <= x <= xmax
        ymin, ymax: float
            Only plot trajectory points for which ymin <= y <= ymax
        zmin, zmax: float
            Only plot trajectory points for which zmin <= z <= zmax
        color: str
            Color to use for the particle trajectories
        line_width: int
            Width of the trajectory lines
        """
        for t in trajectories:
            if not len(t):
                continue
            
            mask = np.full(len(t), True)

            if xmin is not None:
                mask &= t[:, 0] >= xmin
            if xmax is not None:
                mask &= t[:, 0] <= xmax
            if ymin is not None:
                mask &= t[:, 1] >= ymin
            if ymax is not None:
                mask &= t[:, 1] <= ymax
            if zmin is not None:
                mask &= t[:, 2] >= zmin
            if zmax is not None:
                mask &= t[:, 2] <= zmax
            
            t = t[mask]
            
            if not len(t):
                continue
            
            lines = vedo.shapes.Lines(start_pts=t[:-1, :3], end_pts=t[1:, :3], c=color, lw=line_width)
            self.to_plot.append(lines)

    def plot_charge_density(self, excitation, field, color_map='coolwarm'):
        """Plot charge density using the Vedo library.
        
        Parameters
        ---------------------
        excitation: `traceon.excitation.Excitation`
            Excitation applied
        field: `traceon.solver.FieldBEM`
            Field that resulted after solving for the applied excitation
        color_map: str
            Name of the color map to use to color the charge density values
        """
        mesh = excitation.mesh
        
        if not len(mesh.triangles) and not len(mesh.lines):
            raise RuntimeError("Trying to plot empty mesh.")
        
        if len(mesh.triangles):
            meshes = _get_vedo_charge_density_3d(excitation, field, color_map)
            self.to_plot.append(meshes)
            
        if len(mesh.lines):
            lines = _get_vedo_charge_density_2d(excitation, field, color_map)
            self.to_plot.append(lines)
         
        self.is_2d &= mesh.is_2d()
    
    def show(self):
        """Show the figure."""
        plotter = vedo.Plotter() 

        for t in self.to_plot:
            plotter += t
        
        if self.show_legend:
            lb = vedo.LegendBox(self.legend_entries)
            plotter += lb
        
        if self.is_2d:
            plotter.add_global_axes(dict(number_of_divisions=[12, 0, 12], zxgrid=True, xaxis_rotation=90))
        else:
            plotter.add_global_axes(dict(number_of_divisions=[10, 10, 10]))
        
        plotter.look_at(plane='xz')
        plotter.show()

Methods

def plot_charge_density(self, excitation, field, color_map='coolwarm')

Expand source code

def plot_charge_density(self, excitation, field, color_map='coolwarm'):
    """Plot charge density using the Vedo library.
    
    Parameters
    ---------------------
    excitation: `traceon.excitation.Excitation`
        Excitation applied
    field: `traceon.solver.FieldBEM`
        Field that resulted after solving for the applied excitation
    color_map: str
        Name of the color map to use to color the charge density values
    """
    mesh = excitation.mesh
    
    if not len(mesh.triangles) and not len(mesh.lines):
        raise RuntimeError("Trying to plot empty mesh.")
    
    if len(mesh.triangles):
        meshes = _get_vedo_charge_density_3d(excitation, field, color_map)
        self.to_plot.append(meshes)
        
    if len(mesh.lines):
        lines = _get_vedo_charge_density_2d(excitation, field, color_map)
        self.to_plot.append(lines)
     
    self.is_2d &= mesh.is_2d()

Plot charge density using the Vedo library.

Parameters

excitation : Excitation

Excitation applied

field : traceon.solver.FieldBEM

Field that resulted after solving for the applied excitation

color_map : str

Name of the color map to use to color the charge density values

def plot_equipotential_lines(self,
field,
surface,
N0=75,
N1=75,
color_map='coolwarm',
N_isolines=40,
isolines_width=1,
isolines_color='#444444')

Expand source code

def plot_equipotential_lines(self, field, surface, N0=75, N1=75, color_map='coolwarm', N_isolines=40, isolines_width=1, isolines_color='#444444'):
    """Make potential color plot including equipotential lines.

    Parameters
    -------------------------------------
    field: `traceon.solver.Field`
        The field used to compute the potential values (note that any field returned from the solver can be used)
    surface: `traceon.geometry.Surface`
        The surface in 3D space which will be 'colored in'
    N0: int
        Number of pixels to use along the first 'axis' of the surface
    N1: int
        Number of pixels to use along the second 'axis' of the surface
    color_map: str
        Color map to use to color in the surface
    N_isolines: int
        Number of equipotential lines to plot
    isolines_width: int
        The width to use for the isolines. Pass in 0 to disable the isolines.
    isolines_color: str
        Color to use for the isolines"""
    grid = _get_vedo_grid(field, surface, N0, N1)
    isolines = grid.isolines(n=N_isolines).color(isolines_color).lw(isolines_width) # type: ignore
    grid.cmap(color_map)
    self.to_plot.append(grid)
    self.to_plot.append(isolines)

Make potential color plot including equipotential lines.

Parameters

field : traceon.solver.Field

The field used to compute the potential values (note that any field returned from the solver can be used)

surface : Surface

The surface in 3D space which will be 'colored in'

N0 : int

Number of pixels to use along the first 'axis' of the surface

N1 : int

Number of pixels to use along the second 'axis' of the surface

color_map : str

Color map to use to color in the surface

N_isolines : int

Number of equipotential lines to plot

isolines_width : int

The width to use for the isolines. Pass in 0 to disable the isolines.

isolines_color : str

Color to use for the isolines

def plot_mesh(self, mesh, show_normals=False, **colors)

Expand source code

def plot_mesh(self, mesh, show_normals=False, **colors):
    """Plot mesh using the Vedo library. Optionally showing normal vectors.

    Parameters
    ---------------------
    mesh: `traceon.mesher.Mesh`
        The mesh to plot
    show_normals: bool
        Whether to show the normal vectors at every element
    colors: dict of (string, string)
        Use keyword arguments to specify colors, for example `plot_mesh(mesh, lens='blue', ground='green')`
    """
    if not len(mesh.triangles) and not len(mesh.lines):
        raise RuntimeError("Trying to plot empty mesh.")

    triangle_normals, line_normals = None, None
    
    if len(mesh.triangles):
        meshes, triangle_normals = _get_vedo_triangles_and_normals(mesh, **colors)
        self.legend_entries.extend(meshes)
        self.to_plot.append(meshes)
    
    if len(mesh.lines):
        lines, line_normals = _get_vedo_lines_and_normals(mesh, **colors)
        self.legend_entries.extend(lines)
        self.to_plot.append(lines)
     
    if show_normals:
        if triangle_normals is not None:
            self.to_plot.append(triangle_normals)
        if line_normals is not None:
            self.to_plot.append(line_normals)
    
    self.is_2d &= mesh.is_2d()

Plot mesh using the Vedo library. Optionally showing normal vectors.

Parameters

mesh : Mesh

The mesh to plot

show_normals : bool

Whether to show the normal vectors at every element

colors : dict of (string, string)

Use keyword arguments to specify colors, for example plot_mesh(mesh, lens='blue', ground='green')

def plot_trajectories(self,
trajectories,
xmin=None,
xmax=None,
ymin=None,
ymax=None,
zmin=None,
zmax=None,
color='#00AA00',
line_width=1)

Expand source code

def plot_trajectories(self, trajectories, 
            xmin=None, xmax=None,
            ymin=None, ymax=None,
            zmin=None, zmax=None,
            color='#00AA00', line_width=1):
    """Plot particle trajectories.

    Parameters
    ------------------------------------
    trajectories: list of numpy.ndarray
        List of positions as returned by `traceon.tracing.Tracer.__call__`
    xmin, xmax: float
        Only plot trajectory points for which xmin <= x <= xmax
    ymin, ymax: float
        Only plot trajectory points for which ymin <= y <= ymax
    zmin, zmax: float
        Only plot trajectory points for which zmin <= z <= zmax
    color: str
        Color to use for the particle trajectories
    line_width: int
        Width of the trajectory lines
    """
    for t in trajectories:
        if not len(t):
            continue
        
        mask = np.full(len(t), True)

        if xmin is not None:
            mask &= t[:, 0] >= xmin
        if xmax is not None:
            mask &= t[:, 0] <= xmax
        if ymin is not None:
            mask &= t[:, 1] >= ymin
        if ymax is not None:
            mask &= t[:, 1] <= ymax
        if zmin is not None:
            mask &= t[:, 2] >= zmin
        if zmax is not None:
            mask &= t[:, 2] <= zmax
        
        t = t[mask]
        
        if not len(t):
            continue
        
        lines = vedo.shapes.Lines(start_pts=t[:-1, :3], end_pts=t[1:, :3], c=color, lw=line_width)
        self.to_plot.append(lines)

Plot particle trajectories.

Parameters

trajectories : list of numpy.ndarray

List of positions as returned by Tracer.__call__()

xmin, xmax : float

Only plot trajectory points for which xmin <= x <= xmax

ymin, ymax : float

Only plot trajectory points for which ymin <= y <= ymax

zmin, zmax : float

Only plot trajectory points for which zmin <= z <= zmax

color : str

Color to use for the particle trajectories

line_width : int

Width of the trajectory lines

def show(self)

Expand source code

def show(self):
    """Show the figure."""
    plotter = vedo.Plotter() 

    for t in self.to_plot:
        plotter += t
    
    if self.show_legend:
        lb = vedo.LegendBox(self.legend_entries)
        plotter += lb
    
    if self.is_2d:
        plotter.add_global_axes(dict(number_of_divisions=[12, 0, 12], zxgrid=True, xaxis_rotation=90))
    else:
        plotter.add_global_axes(dict(number_of_divisions=[10, 10, 10]))
    
    plotter.look_at(plane='xz')
    plotter.show()

Show the figure.