diff --git a/setup.py b/setup.py index 933429c8..71c8be14 100644 --- a/setup.py +++ b/setup.py @@ -58,6 +58,7 @@ "ds-filter-files = disc_solver.filter_files:main", "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", ], }, cmdclass=versioneer.get_cmdclass(), diff --git a/src/disc_solver/analyse/component_plot.py b/src/disc_solver/analyse/component_plot.py new file mode 100644 index 00000000..689a7285 --- /dev/null +++ b/src/disc_solver/analyse/component_plot.py @@ -0,0 +1,259 @@ +# -*- coding: utf-8 -*- +""" +Plot command for DiscSolver +""" +from numpy import degrees, sqrt, tan +import matplotlib.pyplot as plt + +from ..solve.deriv_funcs import deriv_B_r_func +from ..utils import ODEIndex + +from .utils import ( + single_solution_plotter, analyse_main_wrapper, analysis_func_wrapper, + common_plotting_options, get_common_plot_args, plot_output_wrapper, + AnalysisError, +) + +plt.style.use("bmh") + + +def plot_parser(parser): + """ + Add arguments for plot command to parser + """ + common_plotting_options(parser) + parser.add_argument("--only", default=None) + return parser + + +def get_plot_args(args): + """ + Parse plot args + """ + return { + "only": args.get("only", None) + } + + +@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='-', + stop=90, figargs=None, title=None, close=True, only=None, +): + """ + Plot solution to file + """ + # pylint: disable=too-many-function-args,unexpected-keyword-arg + fig = generate_plot( + soln, soln_range, linestyle=linestyle, figargs=figargs, title=title, + stop=stop, only=only + ) + + return plot_output_wrapper( + fig, file=plot_filename, show=show, close=close + ) + + +@single_solution_plotter +def generate_plot( + soln, *, linestyle='-', stop=90, figargs=None, only=None +): + """ + Generate plot, with enough freedom to be able to format fig + """ + if figargs is None: + figargs = {} + + angles = soln.angles + indexes = degrees(angles) <= stop + plot_angles = degrees(angles[indexes]) + + components = compute_components(soln, indexes, angles[indexes]) + + if only: + fig, ax = plt.subplots(constrained_layout=True, **figargs) + comps = components.get(only) + if comps is None: + raise AnalysisError("{} not found".format(only)) + + for label, values in comps.items(): + ax.plot(plot_angles, values, linestyle, label=label) + + ax.set_xlabel("angle from plane (°)") + ax.set_ylabel(only) + ax.legend(loc=0) + + else: + fig, axes = plt.subplots( + nrows=2, ncols=4, constrained_layout=True, sharex=True, + gridspec_kw=dict(hspace=0), **figargs + ) + + # only add label to bottom plots + for ax in axes[1]: + ax.set_xlabel("angle from plane (°)") + + for ax, comp_pair in zip(axes.flat, components.items()): + eq_name, comps = comp_pair + ax.set_ylabel(eq_name) + for label, values in comps.items(): + ax.plot(plot_angles, values, linestyle, label=label) + + ax.legend(loc=0) + + return fig + + +def compute_components(soln, indexes, angles): + """ + Return mapping containing all the components + """ + solution = soln.solution + cons = soln.initial_conditions + + B_r = solution[indexes, ODEIndex.B_r] + B_φ = solution[indexes, ODEIndex.B_φ] + B_θ = solution[indexes, ODEIndex.B_θ] + v_r = solution[indexes, ODEIndex.v_r] + v_φ = solution[indexes, ODEIndex.v_φ] + v_θ = solution[indexes, ODEIndex.v_θ] + ρ = solution[indexes, ODEIndex.ρ] + deriv_B_φ = solution[indexes, ODEIndex.B_φ_prime] + η_O = solution[indexes, ODEIndex.η_O] + η_A = solution[indexes, ODEIndex.η_A] + η_H = solution[indexes, ODEIndex.η_H] + + γ = cons.γ + a_0 = cons.a_0 + norm_kepler_sq = cons.norm_kepler_sq + + norm_kepler = sqrt(norm_kepler_sq) + + B_mag = sqrt(B_r**2 + B_φ**2 + B_θ**2) + norm_B_r, norm_B_φ, norm_B_θ = B_r/B_mag, B_φ/B_mag, B_θ/B_mag + + deriv_B_r = deriv_B_r_func( + B_r=B_r, B_φ=B_φ, B_θ=B_θ, η_O=η_O, η_H=η_H, η_A=η_A, θ=angles, + v_r=v_r, v_θ=v_θ, deriv_B_φ=deriv_B_φ, γ=γ, + ) + + deriv_v_θ = ( + v_r / 2 * (v_θ ** 2 - 4 * γ) + v_θ * ( + tan(angles) * (v_φ ** 2 + 1) + a_0 / ρ * ( + (1/4 - γ) * B_θ * B_r + B_r * deriv_B_r + + B_φ * deriv_B_φ - B_φ ** 2 * tan(angles) + ) + ) + ) / ((1 - v_θ) * (1 + v_θ)) + + components = { + "radial momentum": { + "radial terms": v_r ** 2 + v_θ ** 2 + 5 / 2 - 2 * γ, + "keplerian terms": (v_φ - norm_kepler) * (v_φ + norm_kepler), + "magnetic terms": ( + a_0 / ρ * (B_θ * deriv_B_r + (1/4 - γ) * (B_θ ** 2 + B_φ ** 2)) + ), + "all terms": ( + v_r ** 2 / 2 + v_θ ** 2 + 5/2 - 2 * γ + + (v_φ - norm_kepler) * (v_φ + norm_kepler) + a_0 / ρ * ( + B_θ * deriv_B_r + (1/4 - γ) * (B_θ ** 2 + B_φ ** 2) + ) + ), + }, + "azimuthal momentum": { + "non-magnetic terms": v_φ * v_θ * tan(angles) - v_φ * v_r / 2, + "magnetic terms": a_0 / ρ * ( + B_θ * deriv_B_φ - (1/4 - γ) * B_r * B_φ - + B_θ * B_φ * tan(angles) + ), + "all terms": ( + v_φ * v_θ * tan(angles) - v_φ * v_r / 2 + a_0 / ρ * ( + B_θ * deriv_B_φ - (1/4 - γ) * B_r * B_φ - + B_θ * B_φ * tan(angles) + ) + ), + }, + "polar momentum": { + "no-B no-v_φ terms": v_r / 2 * (v_θ ** 2 - 4 * γ), + "no-B v_φ terms": v_θ * tan(angles) * (v_φ ** 2 + 1), + "magnetic terms": ( + v_θ * a_0 / ρ * ( + (1/4 - γ) * B_θ * B_r + B_r * deriv_B_r + + B_φ * deriv_B_φ - B_φ ** 2 * tan(angles) + ) + ), + "all terms": ( + v_r / 2 * (v_θ ** 2 - 4 * γ) + v_θ * ( + tan(angles) * (v_φ ** 2 + 1) + a_0 / ρ * ( + (1/4 - γ) * B_θ * B_r + B_r * deriv_B_r + + B_φ * deriv_B_φ - B_φ ** 2 * tan(angles) + ) + ) + ), + }, + "continuity": { + "tan term": v_θ * tan(angles), + "non-tan terms": - (2 * γ * v_r + deriv_v_θ), + "all terms": - (2 * γ * v_r + deriv_v_θ) + v_θ * tan(angles), + }, + "solenoid condition": { + "B_θ terms": B_θ * tan(angles), + "B_r terms": - (γ + 3/4) * B_r, + "all terms": B_θ * tan(angles) - (γ + 3/4) * B_r, + }, + "polar induction": { + "v terms": (v_θ * B_r - v_r * B_θ) / ( + η_O + η_A * (1 - norm_B_φ) * (1 + norm_B_φ) + ), + "no-η terms": - B_θ * (1/4 - γ), + "$B_φ'$ terms": - deriv_B_φ * ( + η_H * norm_B_θ + η_A * norm_B_r * norm_B_φ + ) / ( + η_O + η_A * (1 - norm_B_φ) * (1 + norm_B_φ) + ), + "$B_φ$ terms": B_φ * ( + η_A * norm_B_φ * ( + norm_B_r * tan(angles) - norm_B_θ * (1/4 - γ) + ) + η_H * ( + norm_B_r * (1/4 - γ) + norm_B_θ * tan(angles) + ) + ) / ( + η_O + η_A * (1 - norm_B_φ) * (1 + norm_B_φ) + ), + "all terms": ( + ( + v_θ * B_r - v_r * B_θ - deriv_B_φ * ( + η_H * norm_B_θ + + η_A * norm_B_r * norm_B_φ + ) + B_φ * ( + η_A * norm_B_φ * ( + norm_B_r * tan(angles) - + norm_B_θ * (1/4 - γ) + ) + η_H * ( + norm_B_r * (1/4 - γ) + + norm_B_θ * tan(angles) + ) + ) + ) / ( + η_O + η_A * (1 - norm_B_φ) * (1 + norm_B_φ) + ) - B_θ * (1/4 - γ) + ), + }, + } + + return components