From 8a1671f89225a0b5fd580190e9aee69ccc377e44 Mon Sep 17 00:00:00 2001 From: James Tocknell Date: Thu, 6 Dec 2018 12:54:22 +1100 Subject: [PATCH] Add vert plot --- setup.py | 1 + src/disc_solver/analyse/vert_plot.py | 172 +++++++++++++++++++++++++++ tests/test_run.py | 10 ++ 3 files changed, 183 insertions(+) create mode 100644 src/disc_solver/analyse/vert_plot.py diff --git a/setup.py b/setup.py index 56c1dcba..43a61430 100644 --- a/setup.py +++ b/setup.py @@ -60,6 +60,7 @@ "ds-j-e-plot = disc_solver.analyse.j_e_plot:j_e_plot_main", "ds-plot-taylor-space = disc_solver.solve.taylor_space:main", "ds-component-plot = disc_solver.analyse.component_plot:plot_main", + "ds-vert-plot = disc_solver.analyse.vert_plot:plot_main", ], }, cmdclass=versioneer.get_cmdclass(), diff --git a/src/disc_solver/analyse/vert_plot.py b/src/disc_solver/analyse/vert_plot.py new file mode 100644 index 00000000..f0dd4ead --- /dev/null +++ b/src/disc_solver/analyse/vert_plot.py @@ -0,0 +1,172 @@ +# -*- coding: utf-8 -*- +""" +Plot command for DiscSolver +""" +from numpy import degrees, sqrt, ones as np_ones + +from ..utils import ( + mhd_wave_speeds, MHD_Wave_Index, CylindricalODEIndex, + VERT_MAGNETIC_INDEXES, convert_spherical_to_cylindrical, + get_vertical_scaling, +) + +from .utils import ( + single_solution_plotter, analyse_main_wrapper, analysis_func_wrapper, + common_plotting_options, get_common_plot_args, plot_output_wrapper, + DEFAULT_MPL_STYLE, get_common_arguments, PlotOrdering, +) + + +def plot_parser(parser): + """ + Add arguments for plot command to parser + """ + common_plotting_options(parser) + parser.add_argument("--v_θ", choices=("log", "linear"), default="linear") + parser.add_argument( + "--with-slow", action='store_true', default=False) + parser.add_argument( + "--with-alfven", action='store_true', default=False) + parser.add_argument( + "--with-fast", action='store_true', default=False) + parser.add_argument( + "--with-sonic", action='store_true', default=False) + return parser + + +def get_plot_args(args): + """ + Parse plot args + """ + return { + "with_slow": args.get("with_slow", False), + "with_alfven": args.get("with_alfven", False), + "with_fast": args.get("with_fast", False), + "with_sonic": args.get("with_sonic", False), + } + + +@analyse_main_wrapper( + "Main plotter for DiscSolver", + plot_parser, + cmd_parser_splitters={ + "common_plot_args": get_common_plot_args, + "plot_args": get_plot_args, + } +) +def plot_main(soln, *, soln_range, common_plot_args, plot_args): + """ + Entry point for ds-plot + """ + return plot(soln, soln_range=soln_range, **common_plot_args, **plot_args) + + +@analysis_func_wrapper +def plot( + soln, *, soln_range=None, plot_filename=None, show=False, linestyle='-', + with_slow=False, with_alfven=False, with_fast=False, with_sonic=False, + stop=90, figargs=None, title=None, close=True, filename, + mpl_style=DEFAULT_MPL_STYLE, with_version=True +): + """ + Plot solution to file + """ + # pylint: disable=too-many-function-args,unexpected-keyword-arg + fig = generate_plot( + soln, soln_range, linestyle=linestyle, with_slow=with_slow, + with_alfven=with_alfven, with_fast=with_fast, with_sonic=with_sonic, + stop=stop, figargs=figargs, title=title, filename=filename, + mpl_style=mpl_style, with_version=with_version, + ) + + return plot_output_wrapper( + fig, file=plot_filename, show=show, close=close + ) + + +@single_solution_plotter +def generate_plot( + fig, soln, *, linestyle='-', with_slow=False, with_alfven=False, + with_fast=False, with_sonic=False, stop=90, use_E_r=False +): + """ + Generate plot, with enough freedom to be able to format fig + """ + solution = soln.solution + angles = soln.angles + inp = soln.solution_input + cons = soln.initial_conditions + + heights, vert_soln = convert_spherical_to_cylindrical( + angles, solution, γ=inp.γ, c_s_on_v_k=inp.c_s_on_v_k, use_E_r=use_E_r, + ) + + wave_speeds = sqrt(mhd_wave_speeds( + vert_soln[:, VERT_MAGNETIC_INDEXES], + vert_soln[:, CylindricalODEIndex.ρ], 1, + index=CylindricalODEIndex.B_z, + )) + + indexes = heights <= stop + + ordering = PlotOrdering.E_r_vert if use_E_r else ( + PlotOrdering.B_φ_prime_vert + ) + + param_names = get_common_arguments( + ordering, initial_conditions=cons, v_φ_offset=sqrt( + cons.norm_kepler_sq / get_vertical_scaling( + angles, c_s_on_v_k=inp.c_s_on_v_k + ) + ), + ) + + if with_slow: + param_names[CylindricalODEIndex.v_z]["extras"].append({ + "label": "slow", + "data": wave_speeds[MHD_Wave_Index.slow], + }) + if with_alfven: + param_names[CylindricalODEIndex.v_z]["extras"].append({ + "label": "alfven", + "data": wave_speeds[MHD_Wave_Index.alfven], + }) + if with_fast: + param_names[CylindricalODEIndex.v_z]["extras"].append({ + "label": "fast", + "data": wave_speeds[MHD_Wave_Index.fast], + }) + if with_sonic: + param_names[CylindricalODEIndex.v_z]["extras"].append({ + "label": "sound", + "data": np_ones(len(solution)), + }) + + axes = fig.subplots( + nrows=2, ncols=4, sharex=True, gridspec_kw=dict(hspace=0), + ) + + # only add label to bottom plots + for ax in axes[1]: + ax.set_xlabel("angle from plane (°)") + + axes.shape = len(param_names) + for i, settings in enumerate(param_names): + ax = axes[i] + ax.plot( + degrees(angles[indexes]), + ( + solution[:, i] - settings.get("offset", 0) + )[indexes], linestyle, label=settings["name"] + ) + for extra in settings.get("extras", []): + ax.plot( + degrees(angles[indexes]), + extra["data"][indexes], + label=extra.get("label") + ) + ax.set_ylabel(settings["name"]) + ax.set_yscale(settings.get("scale", "linear")) + if settings.get("legend"): + ax.legend(loc=0) + return fig diff --git a/tests/test_run.py b/tests/test_run.py index 6bdd949f..5a497233 100644 --- a/tests/test_run.py +++ b/tests/test_run.py @@ -14,6 +14,7 @@ from disc_solver.analyse.acc_plot import acc_plot from disc_solver.analyse.validate_plot import validate_plot from disc_solver.analyse.hydro_check_plot import hydro_check_plot +from disc_solver.analyse.vert_plot import plot as vert_plot from disc_solver.analyse.diverge_plot import diverge_main from disc_solver.analyse.conserve_plot import conserve_main from disc_solver.analyse.utils import AnalysisError @@ -443,6 +444,15 @@ def test_j_e_plot_show(self, solution, mpl_interactive): def test_j_e_plot_file(self, solution, plot_file): return j_e_plot(solution, plot_filename=plot_file) + def test_vert_plot_show(self, solution, mpl_interactive): + vert_plot( + solution, show=True, with_slow=True, with_alfven=True, + with_fast=True, with_sonic=True, + ) + + def test_vert_plot_file(self, solution, plot_file): + return vert_plot(solution, plot_filename=plot_file) + def test_taylor_space(mpl_interactive): taylor_space_main()