Add compute_jacobian

Author: aragilar

src/disc_solver/analyse/compute_jacobian.py

@@ -0,0 +1,76 @@
+# -*- 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 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)

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()