From f1240821e60716c18c1404e3867d37d02c73bc85 Mon Sep 17 00:00:00 2001 From: James Tocknell Date: Fri, 10 Jun 2022 13:07:42 +1000 Subject: [PATCH] Add compute_jacobian --- src/disc_solver/analyse/compute_jacobian.py | 78 +++++++++++++++++++++ src/disc_solver/solve/solution.py | 57 +++++++-------- tests/test_run.py | 4 ++ 3 files changed, 111 insertions(+), 28 deletions(-) create mode 100644 src/disc_solver/analyse/compute_jacobian.py diff --git a/src/disc_solver/analyse/compute_jacobian.py b/src/disc_solver/analyse/compute_jacobian.py new file mode 100644 index 00000000..b934e17a --- /dev/null +++ b/src/disc_solver/analyse/compute_jacobian.py @@ -0,0 +1,78 @@ +# -*- coding: utf-8 -*- +""" +""" + +from numpy import zeros + +from ..float_handling import float_type +from ..solve.solution import ode_system +from ..utils import get_solutions +from .utils import ( + analyse_main_wrapper, analysis_func_wrapper, +) + + +def compute_jacobian( + *, γ, a_0, norm_kepler_sq, init_con, θ_scale, η_derivs, use_E_r, θ, params, + eps, +): + 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, check_params=False + ) + + 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)): + 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, + ) + + +@analysis_func_wrapper +def compute_jacobian_from_file( + soln_file, *, soln_range=None, eps, θ_scale=float_type(1), **kwargs +): + + soln = get_solutions(soln_file, soln_range) + return compute_jacobian_from_solution(soln, eps=eps, θ_scale=θ_scale) diff --git a/src/disc_solver/solve/solution.py b/src/disc_solver/solve/solution.py index 650d172f..1635e4b9 100644 --- a/src/disc_solver/solve/solution.py +++ b/src/disc_solver/solve/solution.py @@ -53,7 +53,8 @@ def ode_system( *, γ, a_0, norm_kepler_sq, init_con, θ_scale=float_type(1), with_taylor=False, η_derivs=True, store_internal=True, use_E_r=False, - v_θ_sonic_crit=None, after_sonic=None, deriv_v_θ_func=None + v_θ_sonic_crit=None, after_sonic=None, deriv_v_θ_func=None, + check_params=True, ): """ Set up the system we are solving for. @@ -95,30 +96,30 @@ def rhs_equation(x, params, derivs): """ # pylint: disable=too-many-statements θ = scaled_to_rad(x, θ_scale) - B_r = params[ODEIndex.B_r] - B_φ = params[ODEIndex.B_φ] - B_θ = params[ODEIndex.B_θ] - v_r = params[ODEIndex.v_r] - v_φ = params[ODEIndex.v_φ] - v_θ = params[ODEIndex.v_θ] - ρ = params[ODEIndex.ρ] - η_O = params[ODEIndex.η_O] - η_A = params[ODEIndex.η_A] - η_H = params[ODEIndex.η_H] + B_r = params[..., ODEIndex.B_r] + B_φ = params[..., ODEIndex.B_φ] + B_θ = params[..., ODEIndex.B_θ] + v_r = params[..., ODEIndex.v_r] + v_φ = params[..., ODEIndex.v_φ] + v_θ = params[..., ODEIndex.v_θ] + ρ = params[..., ODEIndex.ρ] + η_O = params[..., ODEIndex.η_O] + η_A = params[..., ODEIndex.η_A] + η_H = params[..., ODEIndex.η_H] if use_E_r: - E_r = params[ODEIndex.E_r] + E_r = params[..., ODEIndex.E_r] B_φ_prime = None else: - B_φ_prime = params[ODEIndex.B_φ_prime] + B_φ_prime = params[..., ODEIndex.B_φ_prime] E_r = None # check sanity of input values - if ρ < 0: + if check_params and ρ < 0: if store_internal: # pylint: disable=unsubscriptable-object problems[θ].append("negative density") return 1 - if v_θ < 0: + if check_params and v_θ < 0: if store_internal: # pylint: disable=unsubscriptable-object problems[θ].append("negative velocity") @@ -242,21 +243,21 @@ def rhs_equation(x, params, derivs): deriv_b_θ=deriv_b_θ, deriv_b_φ=deriv_b_φ, C=C, A=A, ) - derivs[ODEIndex.B_r] = deriv_B_r - derivs[ODEIndex.B_φ] = deriv_B_φ - derivs[ODEIndex.B_θ] = deriv_B_θ - derivs[ODEIndex.v_r] = deriv_v_r - derivs[ODEIndex.v_φ] = deriv_v_φ - derivs[ODEIndex.v_θ] = deriv_v_θ - derivs[ODEIndex.ρ] = deriv_ρ - derivs[ODEIndex.η_O] = deriv_η_O - derivs[ODEIndex.η_A] = deriv_η_A - derivs[ODEIndex.η_H] = deriv_η_H + derivs[..., ODEIndex.B_r] = deriv_B_r + derivs[..., ODEIndex.B_φ] = deriv_B_φ + derivs[..., ODEIndex.B_θ] = deriv_B_θ + derivs[..., ODEIndex.v_r] = deriv_v_r + derivs[..., ODEIndex.v_φ] = deriv_v_φ + derivs[..., ODEIndex.v_θ] = deriv_v_θ + derivs[..., ODEIndex.ρ] = deriv_ρ + derivs[..., ODEIndex.η_O] = deriv_η_O + derivs[..., ODEIndex.η_A] = deriv_η_A + derivs[..., ODEIndex.η_H] = deriv_η_H if use_E_r: - derivs[ODEIndex.E_r] = deriv_E_r + derivs[..., ODEIndex.E_r] = deriv_E_r else: - derivs[ODEIndex.B_φ_prime] = dderiv_B_φ + derivs[..., ODEIndex.B_φ_prime] = dderiv_B_φ if __debug__: log.debug("θ: {}, {}", θ, degrees(θ)) @@ -332,7 +333,7 @@ def jump_across_sonic( store_internal=False, with_taylor=False, θ_scale=θ_scale, use_E_r=use_E_r, ) - derivs = zeros(len(ODEIndex)) + derivs = zeros(len(ODEIndex), dtype=float_type) system(initial_angle, initial_values, derivs) dθ = (jump_before_sonic + jump_after_sonic) / derivs[ODEIndex.v_θ] diff --git a/tests/test_run.py b/tests/test_run.py index 344ef68f..7da68ef8 100644 --- a/tests/test_run.py +++ b/tests/test_run.py @@ -7,6 +7,7 @@ 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 compute_jacobian_from_file 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 @@ -519,6 +520,9 @@ 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_compute_jacobian_from_solution(self, solution): + compute_jacobian_from_file(solution, eps=1e-10) + def test_taylor_space(mpl_interactive): taylor_space_main()