diff --git a/setup.py b/setup.py index 84c854f3..50909ef0 100644 --- a/setup.py +++ b/setup.py @@ -12,7 +12,7 @@ packages = setuptools.find_packages('src'), package_dir = {'': 'src'}, install_requires = [ - "numpy", + "numpy<1.23", "matplotlib>=2.2,<3.4a0", "scikits.odes>=2.3.0dev0", "scipy", @@ -57,6 +57,9 @@ "ds-params-plot = disc_solver.analyse.params_plot:params_main", "ds-taylor-plot = disc_solver.analyse.taylor_plot:taylor_main", "ds-combine-plot = disc_solver.analyse.combine_plot:combine_main", + "ds-jacobian-plot = disc_solver.analyse.jacobian_plot:jacobian_main", + "ds-jacobian-eigenvalue-plot = disc_solver.analyse.compute_jacobian:jacobian_eigenvalues_main", + "ds-jacobian-eigenvector-plot = disc_solver.analyse.compute_jacobian:jacobian_eigenvectors_main", "ds-validate-plot = disc_solver.analyse.validate_plot:validate_plot_main", "ds-hydro-check-plot = disc_solver.analyse.hydro_check_plot:hydro_check_plot_main", "ds-acc-plot = disc_solver.analyse.acc_plot:acc_main", diff --git a/src/disc_solver/analyse/compute_jacobian.py b/src/disc_solver/analyse/compute_jacobian.py new file mode 100644 index 00000000..b9b270db --- /dev/null +++ b/src/disc_solver/analyse/compute_jacobian.py @@ -0,0 +1,338 @@ +# -*- coding: utf-8 -*- +""" +Analysis functions related to the jacobians of existing solutions +""" +from functools import wraps + +from numpy import ( + degrees, zeros, logical_not as npnot, isfinite, nan, full, + array as nparray, +) +from scipy.linalg import eigvals, eig +import matplotlib as mpl +import matplotlib.pyplot as plt + +from ..float_handling import float_type +from ..solve.solution import ode_system +from .utils import ( + single_solution_plotter, common_plotting_options, analyse_main_wrapper, + get_common_plot_args, analysis_func_wrapper, plot_output_wrapper, + DEFAULT_MPL_STYLE, +) + +mpl.rcParams['agg.path.chunksize'] = 1000000 +COLOURS = plt.get_cmap("tab20").colors + + +def plot_parser(parser): + """ + Add arguments for plot command to parser + """ + common_plotting_options(parser) + parser.add_argument( + "--eps", action='store', default=1e-10, type=float + ) + return parser + + +def get_plot_args(args): + """ + Parse plot args + """ + return { + "eps": float(args.get("eps", 1e-10)), + } + + +@analyse_main_wrapper( + "Plot eigenvalues of jacobians for DiscSolver", + plot_parser, + cmd_parser_splitters={ + "common_plot_args": get_common_plot_args, + "plot_args": get_plot_args, + } +) +def jacobian_eigenvalues_main( + soln, *, soln_range, common_plot_args, plot_args +): + """ + Entry point for ds-derivs-plot + """ + return jacobian_eigenvalues_plot( + soln, soln_range=soln_range, **common_plot_args, **plot_args + ) + + +@analyse_main_wrapper( + "Plot eigenvectors of jacobians for DiscSolver", + plot_parser, + cmd_parser_splitters={ + "common_plot_args": get_common_plot_args, + "plot_args": get_plot_args, + } +) +def jacobian_eigenvectors_main( + soln, *, soln_range, common_plot_args, plot_args +): + """ + Entry point for ds-derivs-plot + """ + return jacobian_eigenvectors_plot( + soln, soln_range=soln_range, **common_plot_args, **plot_args + ) + + +def compute_jacobian( + *, γ, a_0, norm_kepler_sq, init_con, θ_scale, η_derivs, use_E_r, θ, params, + eps, +): + """ + Compute jacobian from solution + """ + rhs_eq, _ = ode_system( + γ=γ, a_0=a_0, norm_kepler_sq=norm_kepler_sq, init_con=init_con, + θ_scale=θ_scale, with_taylor=False, η_derivs=η_derivs, + store_internal=False, use_E_r=use_E_r, v_θ_sonic_crit=None, + after_sonic=None, deriv_v_θ_func=None, derivs_post_solution=True, + ) + + solution_length = params.shape[0] + ode_size = params.shape[1] + + J = zeros([solution_length, ode_size, ode_size], dtype=float_type) + + # compute column for each param + for i in range(ode_size): + derivs_h = zeros([solution_length, ode_size], dtype=float_type) + derivs_l = zeros([solution_length, ode_size], dtype=float_type) + params_h = params.copy() + params_l = params.copy() + + # offset only the value associated with this column + params_h[:, i] += eps + params_l[:, i] -= eps + + # we don't check the validity of inputs as we have these from the + # solution + rhs_eq(θ, params_h, derivs_h) + rhs_eq(θ, params_l, derivs_l) + + J[:, :, i] = (derivs_h - derivs_l) / (2 * eps) + return J + + +def compute_jacobian_from_solution(soln, *, eps, θ_scale=float_type(1)): + """ + Compute jacobian from solution + """ + solution = soln.solution + angles = soln.angles + cons = soln.initial_conditions + soln_input = soln.solution_input + + init_con = cons.init_con + γ = cons.γ + a_0 = cons.a_0 + norm_kepler_sq = cons.norm_kepler_sq + + η_derivs = soln_input.η_derivs + use_E_r = soln_input.use_E_r + + return compute_jacobian( + γ=γ, a_0=a_0, norm_kepler_sq=norm_kepler_sq, init_con=init_con, + θ_scale=θ_scale, η_derivs=η_derivs, use_E_r=use_E_r, θ=angles, + params=solution, eps=eps, + ) + + +def jacobian_single_solution_plot_wrapper(func): + """ + Helper wrapper for solutions working with jacobians + """ + @wraps(func) + def jacobian_wrapper( + fig, soln, *args, eps, use_E_r, θ_scale=float_type(1), **kwargs + ): + # pylint: disable=unused-argument + jacobians = compute_jacobian_from_solution( + soln, eps=eps, θ_scale=θ_scale + ) + return func( + fig, *args, jacobians=jacobians, angles=soln.angles, **kwargs + ) + + return jacobian_wrapper + + +@analysis_func_wrapper +def jacobian_eigenvalues_plot( + soln, *, soln_range=None, plot_filename=None, show=False, start=0, stop=90, + figargs=None, linestyle='.', title=None, close=True, filename, + mpl_style=DEFAULT_MPL_STYLE, with_version=True, eps +): + """ + Plot eigenvalues of jacobians of a solution + """ + # pylint: disable=too-many-function-args,unexpected-keyword-arg + # pylint: disable=missing-kwoa + fig = plot_jacobian_eigenvalues( + soln, soln_range, linestyle=linestyle, start=start, stop=stop, + figargs=figargs, title=title, filename=filename, + mpl_style=mpl_style, with_version=with_version, eps=eps + ) + + return plot_output_wrapper( + fig, file=plot_filename, show=show, close=close + ) + + +@analysis_func_wrapper +def jacobian_eigenvectors_plot( + soln, *, soln_range=None, plot_filename=None, show=False, start=0, stop=90, + figargs=None, linestyle='.', title=None, close=True, filename, + mpl_style=DEFAULT_MPL_STYLE, with_version=True, eps +): + """ + Show derivatives + """ + # pylint: disable=too-many-function-args,unexpected-keyword-arg + # pylint: disable=missing-kwoa + fig = plot_jacobian_eigenvectors( + soln, soln_range, linestyle=linestyle, start=start, stop=stop, + figargs=figargs, title=title, filename=filename, + mpl_style=mpl_style, with_version=with_version, eps=eps + ) + + return plot_output_wrapper( + fig, file=plot_filename, show=show, close=close + ) + + +def compute_eigenvalues(jacobian): + """ + Compute eigenvalues of a jacobian matrix + """ + if not isfinite(jacobian).all(): + return full(jacobian.shape[0], nan) + return eigvals(jacobian) + + +def compute_eigenvalues_and_eigenvectors(jacobian): + """ + Compute eigenvalues and eigenvectors of a jacobian matrix + """ + if not isfinite(jacobian).all(): + M = jacobian.shape[0] + return full(M, nan), full([M, M], nan) + return eig(jacobian) + + +def plot_log_complex_data(ax, *, x_data, y_data, **kwargs): + """ + Plot complex data logarithmically + """ + pos_real_data_slice = y_data.real >= 0 + pos_imag_data_slice = y_data.imag >= 0 + ax.plot( + x_data[pos_real_data_slice], + y_data[pos_real_data_slice].real, + marker="^", + linestyle='None', + **kwargs + ) + ax.plot( + x_data[npnot(pos_real_data_slice)], + - y_data[npnot(pos_real_data_slice)].real, + marker="v", + linestyle='None', + **kwargs + ) + ax.plot( + x_data[pos_imag_data_slice], + y_data[pos_imag_data_slice].imag, + marker="2", + linestyle='None', + **kwargs + ) + ax.plot( + x_data[npnot(pos_imag_data_slice)], + - y_data[npnot(pos_imag_data_slice)].imag, + marker="1", + linestyle='None', + **kwargs + ) + + +@single_solution_plotter +@jacobian_single_solution_plot_wrapper +def plot_jacobian_eigenvalues( + fig, *, jacobians, angles, start=0, stop=90, linestyle=',', +): + """ + Generate plot of jacobians + """ + # pylint: disable=unused-argument + indexes = (start <= degrees(angles)) & (degrees(angles) <= stop) + + data = nparray([compute_eigenvalues(j) for j in jacobians]) + + ax = fig.subplots() + ax.set_xlabel("angle from plane (°)") + ax.set_yscale("log") + for i in range(data.shape[1]): + plot_log_complex_data( + ax, x_data=degrees(angles[indexes]), + y_data=data[indexes, i], + color=COLOURS[i], + alpha=.5, + label=str(i) + ) + ax.legend() + return fig + + +@single_solution_plotter +@jacobian_single_solution_plot_wrapper +def plot_jacobian_eigenvectors( + fig, *, jacobians, angles, start=0, stop=90, linestyle=',', +): + """ + Generate plot of jacobians + """ + # pylint: disable=unused-argument + indexes = (start <= degrees(angles)) & (degrees(angles) <= stop) + + eigvalues, eigvecs = zip(*[ + compute_eigenvalues_and_eigenvectors(j) for j in jacobians + ]) + eigvalues = nparray(eigvalues) + eigvecs = nparray(eigvecs) + + axes = fig.subplots( + nrows=2, ncols=6, 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 = 12 + for j in range(eigvalues.shape[1]): + ax = axes[j] + ax.set_yscale("log") + plot_log_complex_data( + ax, x_data=degrees(angles[indexes]), + y_data=eigvalues[indexes, j], + color="black", + alpha=.5, + ) + for i in range(eigvecs.shape[2]): + plot_log_complex_data( + ax, x_data=degrees(angles[indexes]), + y_data=eigvecs[indexes, i, j], + color=COLOURS[i], + alpha=.5, + label=str(i) + ) + ax.legend() + return fig diff --git a/tests/conftest.py b/tests/conftest.py index ab3b2484..11938c8b 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -49,6 +49,13 @@ USE_E_R_SOLUTIONS = [ "mod_hydro_solution_use_E_r", ] +NON_APPROX_SOLUTIONS = [ + "single_solution_default", + "sonic_root_solution_default", + "step_solution_default", + "single_solution_no_internal", + "sonic_root_solution_no_internal", +] @pytest.fixture(scope="session") @@ -213,6 +220,11 @@ def solution_use_E_r(request): return request.getfixturevalue(request.param) +@pytest.fixture(scope="session", params=NON_APPROX_SOLUTIONS) +def solution_not_approx(request): + return request.getfixturevalue(request.param) + + @pytest.fixture() def mpl_interactive(request): import matplotlib as mpl diff --git a/tests/test_run.py b/tests/test_run.py index 344ef68f..39e5d39e 100644 --- a/tests/test_run.py +++ b/tests/test_run.py @@ -7,6 +7,9 @@ from disc_solver.analyse.combine_plot import combine_plot from disc_solver.analyse.compare_plot import compare_plot from disc_solver.analyse.component_plot import plot as component_plot +from disc_solver.analyse.compute_jacobian import ( + jacobian_eigenvalues_plot, jacobian_eigenvectors_plot, +) from disc_solver.analyse.conserve_plot import conserve_main from disc_solver.analyse.derivs_plot import derivs_plot from disc_solver.analyse.diverge_plot import diverge_main @@ -483,12 +486,6 @@ def test_hydro_check_file_no_internal( with pytest.raises(AnalysisError): hydro_check_plot(solution_no_internal, plot_filename=plot_file) - def test_jacobian_show(self, solution, mpl_interactive): - plot(solution, show=True) - - def test_jacobian_file(self, solution, plot_file): - return plot(solution, plot_filename=plot_file) - def test_diverge_show(self, solution, mpl_interactive): diverge_main([solution, '--show']) @@ -519,6 +516,26 @@ def test_vert_plot_file(self, solution, plot_file): def test_stats(self, solution, tmpdir): return stats_main([solution, '--file', str(tmpdir / 'stats.csv')]) + def test_jacobian_eigenvalues_show( + self, solution_not_approx, mpl_interactive + ): + jacobian_eigenvalues_plot(solution_not_approx, show=True, eps=1e-10) + + def test_jacobian_eigenvalues_file(self, solution_not_approx, plot_file): + return jacobian_eigenvalues_plot( + solution_not_approx, plot_filename=plot_file, eps=1e-10 + ) + + def test_jacobian_eigenvectors_show( + self, solution_not_approx, mpl_interactive + ): + jacobian_eigenvectors_plot(solution_not_approx, show=True, eps=1e-10) + + def test_jacobian_eigenvectors_file(self, solution_not_approx, plot_file): + return jacobian_eigenvectors_plot( + solution_not_approx, plot_filename=plot_file, eps=1e-10 + ) + def test_taylor_space(mpl_interactive): taylor_space_main()