Add compute_jacobian

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)

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_θ]

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