Skip to content

Commit

Permalink
some rework on the plottings
Browse files Browse the repository at this point in the history
  • Loading branch information
@MartinPdeS
MartinPdeS committed Oct 24, 2024
1 parent bcda367 commit 02158e3
Show file tree
Hide file tree
Showing 3 changed files with 219 additions and 10 deletions.
211 changes: 209 additions & 2 deletions SuPyMode/helper.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,11 +5,68 @@


def singular_plot_helper(function: Callable) -> Callable:
def wrapper(self, ax: plt.Axes = None, show: bool = True, mode_of_interest: str = 'all', **kwargs) -> plt.Figure:
"""
A decorator that helps in plotting by wrapping a plotting function with additional functionality
such as handling axes creation, setting the figure style, managing legends, and saving figures.
Parameters
----------
function : Callable
The plotting function that is decorated. It should accept `self`, `ax`, and `mode_of_interest`
as parameters.
Returns
-------
Callable
A wrapper function that adds the specified plotting functionalities.
Notes
-----
This decorator expects the decorated function to have the following signature:
`function(self, ax=None, mode_of_interest='all', **kwargs)`.
"""
def wrapper(self, ax: plt.Axes = None, show: bool = True, mode_of_interest: str = 'all', save_filename: str = None, **kwargs) -> plt.Figure:
"""
A wrapped version of the plotting function that provides additional functionality for creating
and managing plots.
Parameters
----------
self : object
The instance of the class calling this method.
ax : plt.Axes, optional
A matplotlib Axes object to draw the plot on. If None, a new figure and axes are created.
Default is None.
show : bool, optional
Whether to display the plot. If False, the plot will not be shown but can still be saved
or returned. Default is True.
mode_of_interest : str, optional
Specifies the mode of interest for the plot. If 'all', all available modes will be plotted.
This parameter is interpreted using the `interpret_mode_of_interest` function. Default is 'all'.
save_filename : str, optional
A file path to save the figure. If None, the figure will not be saved. Default is None.
**kwargs : dict
Additional keyword arguments passed to the decorated function.
Returns
-------
plt.Figure
The matplotlib Figure object created or used for the plot.
Notes
-----
- If no `ax` is provided, a new figure and axes are created using the style context `mps`.
- The legend is only added if there are labels to display.
- If `save_filename` is specified, the figure is saved to the given path.
- The plot is shown if `show` is set to True.
"""
if ax is None:
with plt.style.context(mps):
figure, ax = plt.subplots(1, 1)

else:
figure = ax.get_figure()

mode_of_interest = interpret_mode_of_interest(superset=self, mode_of_interest=mode_of_interest)

function(self, ax=ax, mode_of_interest=mode_of_interest, **kwargs)
Expand All @@ -20,6 +77,9 @@ def wrapper(self, ax: plt.Axes = None, show: bool = True, mode_of_interest: str
if labels:
ax.legend()

if save_filename:
figure.savefig(save_filename)

if show:
plt.show()

Expand All @@ -29,7 +89,65 @@ def wrapper(self, ax: plt.Axes = None, show: bool = True, mode_of_interest: str


def combination_plot_helper(function: Callable) -> Callable:
def wrapper(self, ax: plt.Axes = None, show: bool = True, mode_of_interest: str = 'all', combination: str = 'pairs', **kwargs) -> plt.Figure:
"""
A decorator that enhances a plotting function by adding functionality for handling axes creation,
setting the figure style, interpreting combinations and modes, and managing figure display and saving.
Parameters
----------
function : Callable
The plotting function that is being decorated. It should accept `self`, `ax`, `mode_of_interest`,
and `combination` as parameters.
Returns
-------
Callable
A wrapper function that provides additional functionalities for the plotting function.
Notes
-----
The decorated function should have the following signature:
`function(self, ax=None, mode_of_interest='all', combination='pairs', **kwargs)`.
"""
def wrapper(self, ax: plt.Axes = None, show: bool = True, mode_of_interest: str = 'all', combination: str = 'pairs', save_filename: str = None, **kwargs) -> plt.Figure:
"""
A wrapped version of the plotting function that provides additional functionality for creating
and managing plots with specific combinations of modes.
Parameters
----------
self : object
The instance of the class calling this method.
ax : plt.Axes, optional
A matplotlib Axes object to draw the plot on. If None, a new figure and axes are created.
Default is None.
show : bool, optional
Whether to display the plot. If False, the plot will not be shown but can still be saved
or returned. Default is True.
mode_of_interest : str, optional
Specifies the mode of interest for the plot. If 'all', all available modes will be plotted.
This parameter is interpreted using the `interpret_mode_of_interest` function. Default is 'all'.
combination : str, optional
Specifies the type of combination to plot. The value is interpreted using
`self.interpret_combination` and can be 'pairs', 'triplets', etc. Default is 'pairs'.
save_filename : str, optional
A file path to save the figure. If None, the figure will not be saved. Default is None.
**kwargs : dict
Additional keyword arguments passed to the decorated function.
Returns
-------
plt.Figure
The matplotlib Figure object created or used for the plot.
Notes
-----
- If no `ax` is provided, a new figure and axes are created using the style context `mps`.
- The `mode_of_interest` and `combination` are interpreted based on the instance's methods.
- The legend is only added if there are labels to display.
- If `save_filename` is specified, the figure is saved to the given path.
- The plot is shown if `show` is set to True.
"""
if ax is None:
with plt.style.context(mps):
figure, ax = plt.subplots(1, 1)
Expand All @@ -46,6 +164,9 @@ def wrapper(self, ax: plt.Axes = None, show: bool = True, mode_of_interest: str
if labels:
ax.legend()

if save_filename:
figure.savefig(save_filename)

if show:
plt.show()

Expand All @@ -55,7 +176,48 @@ def wrapper(self, ax: plt.Axes = None, show: bool = True, mode_of_interest: str


def parse_mode_of_interest(plot_function: Callable) -> Callable:
"""
A decorator that parses and interprets the `mode_of_interest` parameter for a given plotting function.
Parameters
----------
plot_function : Callable
The plotting function to be decorated. It should accept `self`, `mode_of_interest`, and other
optional arguments.
Returns
-------
Callable
A wrapper function that interprets the `mode_of_interest` parameter and passes it to the
decorated plotting function.
Notes
-----
The decorated function should have the following signature:
`plot_function(self, *args, mode_of_interest='all', **kwargs)`.
"""

def wrapper(self, *args, mode_of_interest='all', **kwargs):
"""
A wrapped version of the plotting function that interprets the `mode_of_interest` parameter.
Parameters
----------
self : object
The instance of the class calling this method.
*args : tuple
Positional arguments passed to the decorated function.
mode_of_interest : str, optional
Specifies the mode of interest for the plot. This parameter is interpreted using the
`interpret_mode_of_interest` function. Default is 'all'.
**kwargs : dict
Additional keyword arguments passed to the decorated function.
Returns
-------
Any
The return value of the decorated plotting function.
"""
mode_of_interest = interpret_mode_of_interest(
superset=self,
mode_of_interest=mode_of_interest
Expand All @@ -67,7 +229,52 @@ def wrapper(self, *args, mode_of_interest='all', **kwargs):


def parse_combination(plot_function: Callable) -> Callable:
"""
A decorator that parses and interprets the `combination` parameter for a given plotting function,
in addition to the `mode_of_interest` parameter.
Parameters
----------
plot_function : Callable
The plotting function to be decorated. It should accept `self`, `mode_of_interest`, `combination`,
and other optional arguments.
Returns
-------
Callable
A wrapper function that interprets the `mode_of_interest` and `combination` parameters and
passes them to the decorated plotting function.
Notes
-----
The decorated function should have the following signature:
`plot_function(self, *args, mode_of_interest='all', combination='pairs', **kwargs)`.
"""
def wrapper(self, *args, mode_of_interest='all', combination: str = 'pairs', **kwargs):
"""
A wrapped version of the plotting function that interprets the `mode_of_interest` and `combination`
parameters.
Parameters
----------
self : object
The instance of the class calling this method.
*args : tuple
Positional arguments passed to the decorated function.
mode_of_interest : str, optional
Specifies the mode of interest for the plot. This parameter is interpreted using the
`interpret_mode_of_interest` function. Default is 'all'.
combination : str, optional
Specifies the type of combination to plot. The value is interpreted using
`self.interpret_combination` and can be 'pairs', 'triplets', etc. Default is 'pairs'.
**kwargs : dict
Additional keyword arguments passed to the decorated function.
Returns
-------
Any
The return value of the decorated plotting function.
"""
combination = self.interpret_combination(
mode_of_interest=mode_of_interest,
combination=combination
Expand Down
10 changes: 8 additions & 2 deletions SuPyMode/superset_plots.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -183,7 +183,10 @@ def plot_normalized_coupling(
for mode_0, mode_1 in combination:
mode_0.normalized_coupling.plot(ax=ax, other_supermode=mode_1, show=False)

ax.legend()
ax.set(
yscale='linear',
ylim=[0, 10]
)

@combination_plot_helper
def plot_adiabatic(
Expand Down Expand Up @@ -217,7 +220,10 @@ def plot_adiabatic(
for profile in numpy.atleast_1d(add_profile):
profile.render_adiabatic_factor_vs_itr_on_ax(ax=ax, line_style='--')

ax.legend()
ax.set(
yscale='log',
ylim=[1e-5, 1]
)

@parse_mode_of_interest
def plot_field(
Expand Down
8 changes: 2 additions & 6 deletions docs/examples/basic/plot_workflow_02.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@
x_bounds="left", # Mesh x-boundary structure.
y_bounds="bottom", # Mesh y-boundary structure.
boundaries=boundaries, # Set of symmetries to be evaluated, each symmetry add a round of simulation
n_sorted_mode=2, # Total computed and sorted mode.
n_sorted_mode=3, # Total computed and sorted mode.
n_added_mode=2, # Additional computed mode that are not considered later except for field comparison [the higher the better but the slower].
# plot_geometry=True, # Plot the geometry mesh before computation.
debug_mode=0, # Print the iteration step for the solver plus some other important steps.
Expand All @@ -53,10 +53,6 @@

superset = workflow.get_superset()

print(superset.supermodes)

# superset[0].adiabatic.plot(superset[1])

# %%
# Field computation: :math:`E_{i,j}`
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Expand All @@ -75,7 +71,7 @@
# %%
# Adiabatic criterion: :math:`\tilde{C}_{i,j}`
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# _ = superset.plot(plot_type='adiabatic')
_ = superset.plot(plot_type='adiabatic', save_filename='test')


# -

0 comments on commit 02158e3

Please sign in to comment.