Repository: beam-tracing/Scotty (language: python)

Issue:
Launch beam refactoring
(WIP until #74 is merged)

Pulls out the beam launch code into a separate functions and a new file, `scotty.launch`.

Fixes a bug where the toroidal launch angle was not taken into account for finding the launch entry position. This may affect simulations with large toroidal launch angles.

Enables launching from the inboard side (although more changes are required for this to work properly, namely the derivative stencils).

Still to do:
- [x] Add more docs


Discussion / hints:

Fix patch:
diff --git a/docs/api/scotty.density_fit.rst b/docs/api/scotty.density_fit.rst
index bcd3eecd..63b834b8 100644
--- a/docs/api/scotty.density_fit.rst
+++ b/docs/api/scotty.density_fit.rst
@@ -5,4 +5,3 @@ scotty.density\_fit module
    :members:
    :undoc-members:
    :show-inheritance:
-   :special-members: __call__
diff --git a/docs/api/scotty.fun_FFD.rst b/docs/api/scotty.fun_FFD.rst
deleted file mode 100644
index b66fe684..00000000
--- a/docs/api/scotty.fun_FFD.rst
+++ /dev/null
@@ -1,7 +0,0 @@
-scotty.fun\_FFD module
-======================
-
-.. automodule:: scotty.fun_FFD
-   :members:
-   :undoc-members:
-   :show-inheritance:
diff --git a/docs/api/scotty.fun_mix.rst b/docs/api/scotty.fun_mix.rst
deleted file mode 100644
index b6943c0a..00000000
--- a/docs/api/scotty.fun_mix.rst
+++ /dev/null
@@ -1,7 +0,0 @@
-scotty.fun\_mix module
-======================
-
-.. automodule:: scotty.fun_mix
-   :members:
-   :undoc-members:
-   :show-inheritance:
diff --git a/docs/api/scotty.hamiltonian.rst b/docs/api/scotty.hamiltonian.rst
new file mode 100644
index 00000000..304409b2
--- /dev/null
+++ b/docs/api/scotty.hamiltonian.rst
@@ -0,0 +1,7 @@
+scotty.hamiltonian module
+=========================
+
+.. automodule:: scotty.hamiltonian
+   :members:
+   :undoc-members:
+   :show-inheritance:
diff --git a/docs/api/scotty.launch.rst b/docs/api/scotty.launch.rst
new file mode 100644
index 00000000..14fe5f1e
--- /dev/null
+++ b/docs/api/scotty.launch.rst
@@ -0,0 +1,7 @@
+scotty.launch module
+====================
+
+.. automodule:: scotty.launch
+   :members:
+   :undoc-members:
+   :show-inheritance:
diff --git a/docs/api/scotty.rst b/docs/api/scotty.rst
index 576491e6..3d0502ad 100644
--- a/docs/api/scotty.rst
+++ b/docs/api/scotty.rst
@@ -26,13 +26,13 @@ Submodules
    scotty.debugging
    scotty.density_fit
    scotty.fun_CFD
-   scotty.fun_FFD
    scotty.fun_evolution
    scotty.fun_general
-   scotty.fun_mix
    scotty.geometry
+   scotty.hamiltonian
    scotty.hornpy
    scotty.init_bruv
+   scotty.launch
    scotty.lensalot
    scotty.plot3_Torbeam_benchmark
    scotty.plot3_pretty
diff --git a/scotty/beam_me_up.py b/scotty/beam_me_up.py
index 74c335c0..3e86e755 100644
--- a/scotty/beam_me_up.py
+++ b/scotty/beam_me_up.py
@@ -66,7 +66,6 @@
 
 from __future__ import annotations
 import numpy as np
-import math
 from scipy import integrate as integrate
 from scipy import constants as constants
 import matplotlib.pyplot as plt
@@ -81,36 +80,23 @@
     contract_special,
     make_unit_vector_from_cross_product,
     find_x0,
-    find_H,
     find_waist,
-)
-from scotty.fun_general import (
-    find_inverse_2D,
+    freq_GHz_to_angular_frequency,
+    angular_frequency_to_wavenumber,
     find_Psi_3D_lab,
-    find_q_lab_Cartesian,
-    find_K_lab_Cartesian,
-    find_K_lab,
-    find_Psi_3D_lab_Cartesian,
 )
+from scotty.fun_general import find_q_lab_Cartesian, find_Psi_3D_lab_Cartesian
 from scotty.fun_general import find_normalised_plasma_freq, find_normalised_gyro_freq
 from scotty.fun_general import find_epsilon_para, find_epsilon_perp, find_epsilon_g
 from scotty.fun_general import find_dbhat_dR, find_dbhat_dZ
-from scotty.fun_general import (
-    find_d_poloidal_flux_dR,
-    find_d_poloidal_flux_dZ,
-    find_Psi_3D_plasma,
-)
 from scotty.fun_general import find_H_Cardano, find_D
-from scotty.fun_evolution import ray_evolution_2D_fun, beam_evolution_fun
 from scotty.fun_evolution import (
-    find_grad_grad_H_vectorised,
-    find_gradK_grad_H_vectorised,
-    find_gradK_gradK_H_vectorised,
+    ray_evolution_2D_fun,
+    beam_evolution_fun,
+    pack_beam_parameters,
+    unpack_beam_parameters,
 )
 
-from scotty.fun_FFD import find_dH_dR, find_dH_dZ  # \nabla H
-from scotty.fun_CFD import find_dH_dKR, find_dH_dKZ, find_dH_dKzeta  # \nabla_K H
-
 # For find_B if using efit files directly
 from scotty.fun_CFD import find_dpolflux_dR, find_dpolflux_dZ
 from scotty.density_fit import density_fit, DensityFitLike
@@ -122,6 +108,8 @@
     CurvySlabField,
     EFITField,
 )
+from scotty.hamiltonian import Hamiltonian, hessians
+from scotty.launch import launch_beam
 from scotty.torbeam import Torbeam
 from scotty._version import __version__
 
@@ -131,31 +119,31 @@
 
 # Type hints
 from typing import Optional, Union, Sequence
-from scotty.typing import PathLike
+from scotty.typing import PathLike, FloatArray
 
 
 def beam_me_up(
-    poloidal_launch_angle_Torbeam,
-    toroidal_launch_angle_Torbeam,
-    launch_freq_GHz,
-    mode_flag,
-    launch_beam_width,
-    launch_beam_curvature,
-    launch_position,
+    poloidal_launch_angle_Torbeam: float,
+    toroidal_launch_angle_Torbeam: float,
+    launch_freq_GHz: float,
+    mode_flag: int,
+    launch_beam_width: float,
+    launch_beam_curvature: float,
+    launch_position: FloatArray,
     ## keyword arguments begin
-    vacuumLaunch_flag=True,
+    vacuumLaunch_flag: bool = True,
     find_B_method: Union[str, MagneticField] = "torbeam",
     shot=None,
     equil_time=None,
-    vacuum_propagation_flag=False,
-    Psi_BC_flag=False,
-    poloidal_flux_enter=None,
+    vacuum_propagation_flag: bool = False,
+    Psi_BC_flag: bool = False,
+    poloidal_flux_enter: Optional[float] = None,
     ## Finite-difference and solver parameters
-    delta_R=-0.0001,  # in the same units as data_R_coord
-    delta_Z=0.0001,  # in the same units as data_Z_coord
-    delta_K_R=0.1,  # in the same units as K_R
-    delta_K_zeta=0.1,  # in the same units as K_zeta
-    delta_K_Z=0.1,  # in the same units as K_z
+    delta_R: float = -0.0001,  # in the same units as data_R_coord
+    delta_Z: float = 0.0001,  # in the same units as data_Z_coord
+    delta_K_R: float = 0.1,  # in the same units as K_R
+    delta_K_zeta: float = 0.1,  # in the same units as K_zeta
+    delta_K_Z: float = 0.1,  # in the same units as K_z
     interp_order=5,  # For the 2D interpolation functions
     len_tau=102,  # Number of tau_points to output
     rtol=1e-3,  # for solve_ivp of the beam solver
@@ -214,6 +202,46 @@ def beam_me_up(
 
     Parameters
     ==========
+    toroidal_launch_angle_Torbeam: float
+        Toroidal angle of antenna in TORBEAM convention
+    poloidal_launch_angle_Torbeam: float
+        Poloidal angle of antenna in TORBEAM convention
+    mode_flag: int
+        Either ``+/-1``, used to determine which mode branch to use
+    launch_beam_width: float
+        Width of the beam at launch
+    launch_beam_curvature: float
+        Curvatuve of the beam at launch
+    launch_position: FloatArray
+        Position of the antenna in cylindrical coordinates
+    vacuumLaunch_flag: bool
+        If ``True``, launch beam from vacuum, otherwise beam launch
+        position is inside the plasma already
+    vacuum_propagation_flag: bool
+        If ``True``, run solver from the launch position, and don't
+        use analytical vacuum propagation
+    Psi_BC_flag: bool
+        If ``True``, use matching boundary conditions at plasma entry
+        position, otherwise do no special treatment at plasma boundary
+    poloidal_flux_enter: Optional
+        Normalised poloidal flux label of plasma boundary. Required if
+        ``vacuum_propagation_flag`` is ``True``
+    plasmaLaunch_Psi_3D_lab_Cartesian: FloatArray
+        :math:`\Psi` of beam in lab Cartesian coordinates. Required if
+        ``vacuumLaunch_flag`` is ``False``
+    plasmaLaunch_K: FloatArray
+        Wavevector of beam at launch. Required if
+        ``vacuumLaunch_flag`` is ``False``
+    delta_R: float
+        Finite difference spacing to use for ``R``
+    delta_Z: float
+        Finite difference spacing to use for ``Z``
+    delta_K_R: float
+        Finite difference spacing to use for ``K_R``
+    delta_K_zeta: float
+        Finite difference spacing to use for ``K_zeta``
+    delta_K_Z: float
+        Finite difference spacing to use for ``K_Z``
     find_B_method:
         - ``"efitpp"`` uses magnetic field data from efitpp files
           directly
@@ -281,8 +309,8 @@ def beam_me_up(
     check_input(mode_flag)
 
     # Tidying up the input data
-    launch_angular_frequency = 2 * math.pi * 10.0**9 * launch_freq_GHz
-    wavenumber_K0 = launch_angular_frequency / constants.c
+    launch_angular_frequency = freq_GHz_to_angular_frequency(launch_freq_GHz)
+    wavenumber_K0 = angular_frequency_to_wavenumber(launch_angular_frequency)
 
     # Ensure paths are `pathlib.Path`
     ne_data_path = pathlib.Path(ne_data_path)
@@ -328,388 +356,77 @@ def beam_me_up(
         delta_Z,
     )
 
+    hamiltonian = Hamiltonian(
+        field,
+        launch_angular_frequency,
+        mode_flag,
+        find_density_1D,
+        delta_R,
+        delta_Z,
+        delta_K_R,
+        delta_K_zeta,
+        delta_K_Z,
+    )
+
     # -------------------
     # Launch parameters
     # -------------------
     if vacuumLaunch_flag:
         print("Beam launched from outside the plasma")
-
-        K_R_launch = (
-            -wavenumber_K0
-            * np.cos(toroidal_launch_angle_Torbeam / 180.0 * math.pi)
-            * np.cos(poloidal_launch_angle_Torbeam / 180.0 * math.pi)
-        )  # K_R
-        K_zeta_launch = (
-            -wavenumber_K0
-            * np.sin(toroidal_launch_angle_Torbeam / 180.0 * math.pi)
-            * np.cos(poloidal_launch_angle_Torbeam / 180.0 * math.pi)
-            * launch_position[0]
-        )  # K_zeta
-        K_Z_launch = -wavenumber_K0 * np.sin(
-            poloidal_launch_angle_Torbeam / 180.0 * math.pi
-        )  # K_z
-        launch_K = np.array([K_R_launch, K_zeta_launch, K_Z_launch])
-
-        poloidal_rotation_angle = (90.0 + poloidal_launch_angle_Torbeam) / 180.0 * np.pi
-
-        Psi_w_beam_launch_cartersian = np.array(
-            [
-                [
-                    wavenumber_K0 * launch_beam_curvature
-                    + 2j * launch_beam_width ** (-2),
-                    0,
-                ],
-                [
-                    0,
-                    wavenumber_K0 * launch_beam_curvature
-                    + 2j * launch_beam_width ** (-2),
-                ],
-            ]
-        )
-        identity_matrix_2D = np.array([[1, 0], [0, 1]])
-
-        rotation_matrix_pol = np.array(
-            [
-                [np.cos(poloidal_rotation_angle), 0, np.sin(poloidal_rotation_angle)],
-                [0, 1, 0],
-                [-np.sin(poloidal_rotation_angle), 0, np.cos(poloidal_rotation_angle)],
-            ]
-        )
-
-        rotation_matrix_tor = np.array(
-            [
-                [
-                    np.cos(toroidal_launch_angle_Torbeam / 180.0 * math.pi),
-                    np.sin(toroidal_launch_angle_Torbeam / 180.0 * math.pi),
-                    0,
-                ],
-                [
-                    -np.sin(toroidal_launch_angle_Torbeam / 180.0 * math.pi),
-                    np.cos(toroidal_launch_angle_Torbeam / 180.0 * math.pi),
-                    0,
-                ],
-                [0, 0, 1],
-            ]
-        )
-
-        rotation_matrix = np.matmul(rotation_matrix_pol, rotation_matrix_tor)
-        rotation_matrix_inverse = np.transpose(rotation_matrix)
-
-        Psi_3D_beam_launch_cartersian = np.array(
-            [
-                [
-                    Psi_w_beam_launch_cartersian[0][0],
-                    Psi_w_beam_launch_cartersian[0][1],
-                    0,
-                ],
-                [
-                    Psi_w_beam_launch_cartersian[1][0],
-                    Psi_w_beam_launch_cartersian[1][1],
-                    0,
-                ],
-                [0, 0, 0],
-            ]
-        )
-        Psi_3D_lab_launch_cartersian = np.matmul(
-            rotation_matrix_inverse,
-            np.matmul(Psi_3D_beam_launch_cartersian, rotation_matrix),
-        )
-        Psi_3D_lab_launch = find_Psi_3D_lab(
-            Psi_3D_lab_launch_cartersian,
-            launch_position[0],
-            launch_position[1],
-            K_R_launch,
-            K_zeta_launch,
+        (
+            K_initial,
+            initial_position,
+            launch_K,
+            Psi_3D_lab_initial,
+            Psi_3D_lab_launch,
+            Psi_3D_lab_entry,
+            Psi_3D_lab_entry_cartersian,
+            distance_from_launch_to_entry,
+        ) = launch_beam(
+            toroidal_launch_angle_Torbeam=toroidal_launch_angle_Torbeam,
+            poloidal_launch_angle_Torbeam=poloidal_launch_angle_Torbeam,
+            launch_beam_width=launch_beam_width,
+            launch_beam_curvature=launch_beam_curvature,
+            launch_position=launch_position,
+            launch_angular_frequency=launch_angular_frequency,
+            field=field,
+            hamiltonian=hamiltonian,
+            vacuum_propagation_flag=vacuum_propagation_flag,
+            Psi_BC_flag=Psi_BC_flag,
+            poloidal_flux_enter=poloidal_flux_enter,
+            delta_R=delta_R,
+            delta_Z=delta_Z,
         )
-
-        if vacuum_propagation_flag:
-            Psi_w_beam_inverse_launch_cartersian = find_inverse_2D(
-                Psi_w_beam_launch_cartersian
-            )
-
-            # Finds entry point
-            search_Z_end = launch_position[2] - launch_position[0] * np.tan(
-                np.radians(poloidal_launch_angle_Torbeam)
-            )
-            numberOfCoarseSearchPoints = 50
-            R_coarse_search_array = np.linspace(
-                launch_position[0], 0, numberOfCoarseSearchPoints
-            )
-            Z_coarse_search_array = np.linspace(
-                launch_position[2], search_Z_end, numberOfCoarseSearchPoints
-            )
-            poloidal_flux_coarse_search_array = np.zeros(numberOfCoarseSearchPoints)
-            poloidal_flux_coarse_search_array = field.poloidal_flux(
-                R_coarse_search_array, Z_coarse_search_array
-            )
-            meets_flux_condition_array = (
-                poloidal_flux_coarse_search_array < 0.9 * poloidal_flux_enter
-            )
-            dummy_array = np.array(range(numberOfCoarseSearchPoints))
-            indices_inside_for_sure_array = dummy_array[meets_flux_condition_array]
-            first_inside_index = indices_inside_for_sure_array[0]
-            numberOfFineSearchPoints = 7500
-            R_fine_search_array = np.linspace(
-                launch_position[0],
-                R_coarse_search_array[first_inside_index],
-                numberOfFineSearchPoints,
-            )
-            Z_fine_search_array = np.linspace(
-                launch_position[2],
-                Z_coarse_search_array[first_inside_index],
-                numberOfFineSearchPoints,
-            )
-            poloidal_fine_search_array = np.zeros(numberOfFineSearchPoints)
-            poloidal_fine_search_array = field.poloidal_flux(
-                R_fine_search_array, Z_fine_search_array
-            )
-            entry_index = find_nearest(poloidal_fine_search_array, poloidal_flux_enter)
-            if poloidal_fine_search_array[entry_index] > poloidal_flux_enter:
-                # The first point needs to be in the plasma
-                # If the first point is outside, then there will be errors when the gradients are calculated
-                entry_index = entry_index + 1
-            entry_position = np.zeros(3)  # R,Z
-            entry_position[0] = R_fine_search_array[entry_index]
-            entry_position[1] = (
-                K_zeta_launch
-                / K_R_launch
-                * (1 / launch_position[0] - 1 / entry_position[0])
-            )
-            entry_position[2] = Z_fine_search_array[entry_index]
-
-            distance_from_launch_to_entry = np.sqrt(
-                launch_position[0] ** 2
-                + entry_position[0] ** 2
-                - 2
-                * launch_position[0]
-                * entry_position[0]
-                * np.cos(entry_position[1] - launch_position[1])
-                + (launch_position[2] - entry_position[2]) ** 2
-            )
-            # Entry point found
-
-            # Calculate entry parameters from launch parameters
-            # That is, find beam at start of plasma given its parameters at the antenna
-            K_lab_launch = np.array([K_R_launch, K_zeta_launch, K_Z_launch])
-            K_lab_Cartesian_launch = find_K_lab_Cartesian(K_lab_launch, launch_position)
-            K_lab_Cartesian_entry = K_lab_Cartesian_launch
-            entry_position_Cartesian = find_q_lab_Cartesian(entry_position)
-            K_lab_entry = find_K_lab(K_lab_Cartesian_entry, entry_position_Cartesian)
-
-            K_R_entry = K_lab_entry[0]  # K_R
-            K_zeta_entry = K_lab_entry[1]
-            K_Z_entry = K_lab_entry[2]  # K_z
-
-            Psi_w_beam_inverse_entry_cartersian = (
-                distance_from_launch_to_entry / (wavenumber_K0) * identity_matrix_2D
-                + Psi_w_beam_inverse_launch_cartersian
-            )
-            Psi_w_beam_entry_cartersian = find_inverse_2D(
-                Psi_w_beam_inverse_entry_cartersian
-            )
-
-            Psi_3D_beam_entry_cartersian = np.array(
-                [
-                    [
-                        Psi_w_beam_entry_cartersian[0][0],
-                        Psi_w_beam_entry_cartersian[0][1],
-                        0,
-                    ],
-                    [
-                        Psi_w_beam_entry_cartersian[1][0],
-                        Psi_w_beam_entry_cartersian[1][1],
-                        0,
-                    ],
-                    [0, 0, 0],
-                ]
-            )  # 'entry' is still in vacuum, so the components of Psi along g are all 0 (since \nabla H = 0)
-
-            Psi_3D_lab_entry_cartersian = np.matmul(
-                rotation_matrix_inverse,
-                np.matmul(Psi_3D_beam_entry_cartersian, rotation_matrix),
-            )
-
-            # Convert to cylindrical coordinates
-            Psi_3D_lab_entry = find_Psi_3D_lab(
-                Psi_3D_lab_entry_cartersian,
-                entry_position[0],
-                entry_position[1],
-                K_R_entry,
-                K_zeta_entry,
-            )
-
-            # -------------------
-            # Find initial parameters in plasma
-            # -------------------
-            K_R_initial = K_R_entry
-            K_zeta_initial = K_zeta_entry
-            K_Z_initial = K_Z_entry
-            initial_position = entry_position
-            if Psi_BC_flag:  # Use BCs
-                dH_dKR_initial = find_dH_dKR(
-                    initial_position[0],
-                    initial_position[2],
-                    K_R_initial,
-                    K_zeta_initial,
-                    K_Z_initial,
-                    launch_angular_frequency,
-                    mode_flag,
-                    delta_K_R,
-                    field.poloidal_flux,
-                    find_density_1D,
-                    field.B_R,
-                    field.B_T,
-                    field.B_Z,
-                )
-                dH_dKzeta_initial = find_dH_dKzeta(
-                    initial_position[0],
-                    initial_position[2],
-                    K_R_initial,
-                    K_zeta_initial,
-                    K_Z_initial,
-                    launch_angular_frequency,
-                    mode_flag,
-                    delta_K_zeta,
-                    field.poloidal_flux,
-                    find_density_1D,
-                    field.B_R,
-                    field.B_T,
-                    field.B_Z,
-                )
-                dH_dKZ_initial = find_dH_dKZ(
-                    initial_position[0],
-                    initial_position[2],
-                    K_R_initial,
-                    K_zeta_initial,
-                    K_Z_initial,
-                    launch_angular_frequency,
-                    mode_flag,
-                    delta_K_Z,
-                    field.poloidal_flux,
-                    find_density_1D,
-                    field.B_R,
-                    field.B_T,
-                    field.B_Z,
-                )
-                dH_dR_initial = find_dH_dR(
-                    initial_position[0],
-                    initial_position[2],
-                    K_R_initial,
-                    K_zeta_initial,
-                    K_Z_initial,
-                    launch_angular_frequency,
-                    mode_flag,
-                    delta_R,
-                    field.poloidal_flux,
-                    find_density_1D,
-                    field.B_R,
-                    field.B_T,
-                    field.B_Z,
-                )
-                dH_dZ_initial = find_dH_dZ(
-                    initial_position[0],
-                    initial_position[2],
-                    K_R_initial,
-                    K_zeta_initial,
-                    K_Z_initial,
-                    launch_angular_frequency,
-                    mode_flag,
-                    delta_Z,
-                    field.poloidal_flux,
-                    find_density_1D,
-                    field.B_R,
-                    field.B_T,
-                    field.B_Z,
-                )
-                d_poloidal_flux_d_R_boundary = find_d_poloidal_flux_dR(
-                    initial_position[0],
-                    initial_position[2],
-                    delta_R,
-                    field.poloidal_flux,
-                )
-                d_poloidal_flux_d_Z_boundary = find_d_poloidal_flux_dZ(
-                    initial_position[0],
-                    initial_position[2],
-                    delta_R,
-                    field.poloidal_flux,
-                )
-
-                Psi_3D_lab_initial = find_Psi_3D_plasma(
-                    Psi_3D_lab_entry,
-                    dH_dKR_initial,
-                    dH_dKzeta_initial,
-                    dH_dKZ_initial,
-                    dH_dR_initial,
-                    dH_dZ_initial,
-                    d_poloidal_flux_d_R_boundary.item(),
-                    d_poloidal_flux_d_Z_boundary.item(),
-                )  # . item() to convert variable from type ndarray to float, such that the array elements all have the same type
-            else:  # Do not use BCs
-                Psi_3D_lab_initial = Psi_3D_lab_entry
-
-        else:  # Run solver from the launch position, no analytical vacuum propagation
-            Psi_3D_lab_initial = Psi_3D_lab_launch
-            K_R_initial = K_R_launch
-            K_zeta_initial = K_zeta_launch
-            K_Z_initial = K_Z_launch
-            initial_position = launch_position
-
-            Psi_3D_lab_entry = None
-            distance_from_launch_to_entry = None
-            Psi_3D_lab_entry_cartersian = np.full_like(
-                Psi_3D_lab_launch, fill_value=np.nan
-            )
-
     else:
         print("Beam launched from inside the plasma")
-        K_R_launch = plasmaLaunch_K[0]
-        K_zeta_launch = plasmaLaunch_K[1]
-        K_Z_launch = plasmaLaunch_K[2]
-
         Psi_3D_lab_initial = find_Psi_3D_lab(
             plasmaLaunch_Psi_3D_lab_Cartesian,
             launch_position[0],
+            launch_position[1],
             plasmaLaunch_K[0],
             plasmaLaunch_K[1],
         )
-        K_R_initial = K_R_launch
-        K_zeta_initial = K_zeta_launch
-        K_Z_initial = K_Z_launch
+        K_initial = plasmaLaunch_K
         initial_position = launch_position
+        launch_K = None
+        Psi_3D_lab_launch = None
+        Psi_3D_lab_entry = None
+        Psi_3D_lab_entry_cartersian = None
+        distance_from_launch_to_entry = None
 
-    #        print(K_R_initial)
-    #        print(K_zeta_launch)
-    #        print(K_Z_launch)
-    # -------------------
+    K_R_initial, K_zeta_initial, K_Z_initial = K_initial
 
     # -------------------
 
     # Initial conditions for the solver
-    beam_parameters_initial = np.zeros(17)
-    # This used to be complex, with a length of 11, but the solver throws a warning saying that something is casted to real
-    # It seems to be fine, bu
-
-    beam_parameters_initial[0] = initial_position[0]  # q_R
-    # q_zeta (This should be = 0)
-    beam_parameters_initial[1] = initial_position[1]
-    beam_parameters_initial[2] = initial_position[2]  # q_Z
-
-    beam_parameters_initial[3] = K_R_initial
-    beam_parameters_initial[4] = K_Z_initial
-
-    beam_parameters_initial[5] = np.real(Psi_3D_lab_initial[0, 0])  # Psi_RR
-    beam_parameters_initial[6] = np.real(Psi_3D_lab_initial[1, 1])  # Psi_zetazeta
-    beam_parameters_initial[7] = np.real(Psi_3D_lab_initial[2, 2])  # Psi_ZZ
-    beam_parameters_initial[8] = np.real(Psi_3D_lab_initial[0, 1])  # Psi_Rzeta
-    beam_parameters_initial[9] = np.real(Psi_3D_lab_initial[0, 2])  # Psi_RZ
-    beam_parameters_initial[10] = np.real(Psi_3D_lab_initial[1, 2])  # Psi_zetaZ
-
-    beam_parameters_initial[11] = np.imag(Psi_3D_lab_initial[0, 0])  # Psi_RR
-    beam_parameters_initial[12] = np.imag(Psi_3D_lab_initial[1, 1])  # Psi_zetazeta
-    beam_parameters_initial[13] = np.imag(Psi_3D_lab_initial[2, 2])  # Psi_ZZ
-    beam_parameters_initial[14] = np.imag(Psi_3D_lab_initial[0, 1])  # Psi_Rzeta
-    beam_parameters_initial[15] = np.imag(Psi_3D_lab_initial[0, 2])  # Psi_RZ
-    beam_parameters_initial[16] = np.imag(Psi_3D_lab_initial[1, 2])  # Psi_zetaZ
+    beam_parameters_initial = pack_beam_parameters(
+        initial_position[0],
+        initial_position[1],
+        initial_position[2],
+        K_R_initial,
+        K_Z_initial,
+        Psi_3D_lab_initial,
+    )
 
     ray_parameters_initial = np.real(beam_parameters_initial[0:5])
     ray_parameters_2D_initial = np.delete(ray_parameters_initial, 1)  # Remove q_zeta
@@ -721,18 +438,7 @@ def event_leave_plasma(
         tau,
         ray_parameters_2D,
         K_zeta,
-        launch_angular_frequency,
-        mode_flag,
-        delta_R,
-        delta_Z,
-        delta_K_R,
-        delta_K_zeta,
-        delta_K_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
+        hamiltonian: Hamiltonian,
         poloidal_flux_leave=poloidal_flux_enter,
     ):  # Leave at the same poloidal flux of entry
         q_R = ray_parameters_2D[0]
@@ -754,18 +460,7 @@ def event_leave_LCFS(
         tau,
         ray_parameters_2D,
         K_zeta,
-        launch_angular_frequency,
-        mode_flag,
-        delta_R,
-        delta_Z,
-        delta_K_R,
-        delta_K_zeta,
-        delta_K_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
+        hamiltonian: Hamiltonian,
         poloidal_flux_LCFS=1.0,
     ):
         q_R = ray_parameters_2D[0]
@@ -787,18 +482,7 @@ def event_leave_simulation(
         tau,
         ray_parameters_2D,
         K_zeta,
-        launch_angular_frequency,
-        mode_flag,
-        delta_R,
-        delta_Z,
-        delta_K_R,
-        delta_K_zeta,
-        delta_K_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
+        hamiltonian: Hamiltonian,
         data_R_coord_min=field.R_coord.min(),
         data_R_coord_max=field.R_coord.max(),
         data_Z_coord_min=field.Z_coord.min(),
@@ -822,23 +506,7 @@ def event_leave_simulation(
     # negative value, when function result goes from positive to negative
     event_leave_simulation.direction = -1.0
 
-    def event_cross_resonance(
-        tau,
-        ray_parameters_2D,
-        K_zeta,
-        launch_angular_frequency,
-        mode_flag,
-        delta_R,
-        delta_Z,
-        delta_K_R,
-        delta_K_zeta,
-        delta_K_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    ):
+    def event_cross_resonance(tau, ray_parameters_2D, K_zeta, hamiltonian: Hamiltonian):
         # Currently only works when crossing resonance.
         # To implement crossing of higher harmonics as well
         delta_gyro_freq = 0.01
@@ -861,23 +529,7 @@ def event_cross_resonance(
     # Stop the solver when the beam is in a region of resonance
     event_cross_resonance.terminal = True
 
-    def event_reach_cutoff(
-        tau,
-        ray_parameters_2D,
-        K_zeta,
-        launch_angular_frequency,
-        mode_flag,
-        delta_R,
-        delta_Z,
-        delta_K_R,
-        delta_K_zeta,
-        delta_K_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    ):
+    def event_reach_cutoff(tau, ray_parameters_2D, K_zeta, hamiltonian: Hamiltonian):
         # To find tau of the cut-off, that is the location where the
         # wavenumber K is minimised
         # This function finds the turning points
@@ -889,53 +541,11 @@ def event_reach_cutoff(
         K_Z = ray_parameters_2D[3]
         K_magnitude = np.sqrt(K_R**2 + K_Z**2 + K_zeta**2 / q_R**2)
 
-        dH_dKR = find_dH_dKR(
-            q_R,
-            q_Z,
-            K_R,
-            K_zeta,
-            K_Z,
-            launch_angular_frequency,
-            mode_flag,
-            delta_K_R,
-            interp_poloidal_flux,
-            find_density_1D,
-            find_B_R,
-            find_B_T,
-            find_B_Z,
-        )
-        dH_dR = find_dH_dR(
-            q_R,
-            q_Z,
-            K_R,
-            K_zeta,
-            K_Z,
-            launch_angular_frequency,
-            mode_flag,
-            delta_R,
-            interp_poloidal_flux,
-            find_density_1D,
-            find_B_R,
-            find_B_T,
-            find_B_Z,
-        )
-        dH_dZ = find_dH_dZ(
-            q_R,
-            q_Z,
-            K_R,
-            K_zeta,
-            K_Z,
-            launch_angular_frequency,
-            mode_flag,
-            delta_Z,
-            interp_poloidal_flux,
-            find_density_1D,
-            find_B_R,
-            find_B_T,
-            find_B_Z,
-        )
+        dH = hamiltonian.derivatives(q_R, q_Z, K_R, K_zeta, K_Z)
 
-        d_K_d_tau = -(1 / K_magnitude) * (dH_dR * K_R + dH_dZ * K_Z + dH_dKR * q_R)
+        d_K_d_tau = -(1 / K_magnitude) * (
+            dH["dH_dR"] * K_R + dH["dH_dZ"] * K_Z + dH["dH_dKR"] * q_R
+        )
         return d_K_d_tau
 
     event_reach_cutoff.direction = 1.0
@@ -950,21 +560,7 @@ def event_reach_cutoff(
     # If the ray hasn't left the plasma by the time this tau is reached, the solver gives up
     tau_max = 10**5
     # Stuff the solver needs to evolve beam_parameters
-    solver_arguments = (
-        K_zeta_initial,
-        launch_angular_frequency,
-        mode_flag,
-        delta_R,
-        delta_Z,
-        delta_K_R,
-        delta_K_zeta,
-        delta_K_Z,
-        field.poloidal_flux,
-        find_density_1D,
-        field.B_R,
-        field.B_T,
-        field.B_Z,
-    )  # Stuff the solver needs to evolve beam_parameters
+    solver_arguments = (K_zeta_initial, hamiltonian)
 
     """
     Propagate a ray. Quickly finds tau at which the ray leaves the plasma, as well as estimates location of cut-off
@@ -1154,7 +750,7 @@ def event_reach_cutoff(
         - suffix _fine for variables in this section
         - I guess I could've used 'dense_output' for the section above and got the information from there, and one day I'll go and see which method works better/faster
         """
-        max_tau_idx = np.argmax(tau_ray[tau_ray <= tau_leave])
+        max_tau_idx = int(np.argmax(tau_ray[tau_ray <= tau_leave]))
 
         K_magnitude_ray = (
             (np.real(ray_parameters_2D[2, :max_tau_idx])) ** 2
@@ -1224,7 +820,12 @@ def event_reach_cutoff(
     )
 
     solver_end_time = time.time()
-    print("Time taken (beam solver)", solver_end_time - solver_start_time, "s")
+    solver_time = solver_end_time - solver_start_time
+    print(f"Time taken (beam solver) {solver_time}s")
+    print(f"Number of beam evolution evaluations: {solver_beam_output.nfev}")
+    print(
+        f"Time per beam evolution evaluation: {solver_time / solver_beam_output.nfev}"
+    )
 
     beam_parameters = solver_beam_output.y
     tau_array = solver_beam_output.t
@@ -1232,35 +833,14 @@ def event_reach_cutoff(
 
     numberOfDataPoints = len(tau_array)
 
-    q_R_array = np.real(beam_parameters[0, :])
-    q_zeta_array = np.real(beam_parameters[1, :])
-    q_Z_array = np.real(beam_parameters[2, :])
-
-    K_R_array = np.real(beam_parameters[3, :])
-    K_Z_array = np.real(beam_parameters[4, :])
-
-    Psi_3D_output = np.zeros([numberOfDataPoints, 3, 3], dtype="complex128")
-    Psi_3D_output[:, 0, 0] = (
-        beam_parameters[5, :] + 1j * beam_parameters[11, :]
-    )  # d (Psi_RR) / d tau
-    Psi_3D_output[:, 1, 1] = (
-        beam_parameters[6, :] + 1j * beam_parameters[12, :]
-    )  # d (Psi_zetazeta) / d tau
-    Psi_3D_output[:, 2, 2] = (
-        beam_parameters[7, :] + 1j * beam_parameters[13, :]
-    )  # d (Psi_ZZ) / d tau
-    Psi_3D_output[:, 0, 1] = (
-        beam_parameters[8, :] + 1j * beam_parameters[14, :]
-    )  # d (Psi_Rzeta) / d tau
-    Psi_3D_output[:, 0, 2] = (
-        beam_parameters[9, :] + 1j * beam_parameters[15, :]
-    )  # d (Psi_RZ) / d tau
-    Psi_3D_output[:, 1, 2] = (
-        beam_parameters[10, :] + 1j * beam_parameters[16, :]
-    )  # d (Psi_zetaZ) / d tau
-    Psi_3D_output[:, 1, 0] = Psi_3D_output[:, 0, 1]
-    Psi_3D_output[:, 2, 0] = Psi_3D_output[:, 0, 2]
-    Psi_3D_output[:, 2, 1] = Psi_3D_output[:, 1, 2]
+    (
+        q_R_array,
+        q_zeta_array,
+        q_Z_array,
+        K_R_array,
+        K_Z_array,
+        Psi_3D_output,
+    ) = unpack_beam_parameters(beam_parameters)
 
     print("Main loop complete")
     # -------------------
@@ -1319,53 +899,17 @@ def event_reach_cutoff(
     poloidal_flux_output = field.poloidal_flux(q_R_array, q_Z_array)
     electron_density_output = np.asfarray(find_density_1D(poloidal_flux_output))
 
-    # Calculates nabla_K H
-    dH_dKR_output = find_dH_dKR(
-        q_R_array,
-        q_Z_array,
-        K_R_array,
-        K_zeta_initial,
-        K_Z_array,
-        launch_angular_frequency,
-        mode_flag,
-        delta_K_R,
-        field.poloidal_flux,
-        find_density_1D,
-        field.B_R,
-        field.B_T,
-        field.B_Z,
-    )
-    dH_dKzeta_output = find_dH_dKzeta(
-        q_R_array,
-        q_Z_array,
-        K_R_array,
-        K_zeta_initial,
-        K_Z_array,
-        launch_angular_frequency,
-        mode_flag,
-        delta_K_zeta,
-        field.poloidal_flux,
-        find_density_1D,
-        field.B_R,
-        field.B_T,
-        field.B_Z,
-    )
-    dH_dKZ_output = find_dH_dKZ(
-        q_R_array,
-        q_Z_array,
-        K_R_array,
-        K_zeta_initial,
-        K_Z_array,
-        launch_angular_frequency,
-        mode_flag,
-        delta_K_Z,
-        field.poloidal_flux,
-        find_density_1D,
-        field.B_R,
-        field.B_T,
-        field.B_Z,
+    dH = hamiltonian.derivatives(
+        q_R_array, q_Z_array, K_R_array, K_zeta_initial, K_Z_array, second_order=True
     )
 
+    dH_dR_output = dH["dH_dR"]
+    dH_dZ_output = dH["dH_dZ"]
+    dH_dKR_output = dH["dH_dKR"]
+    dH_dKzeta_output = dH["dH_dKzeta"]
+    dH_dKZ_output = dH["dH_dKZ"]
+
+    # Calculates nabla_K H
     # Calculates g_hat
     g_hat_output = np.zeros([numberOfDataPoints, 3])
     g_magnitude_output = (
@@ -1426,66 +970,19 @@ def event_reach_cutoff(
     # But good for checking whether things are working properly
     # -------------------
     #
-    H_output = find_H(
-        q_R_array,
-        q_Z_array,
-        K_R_array,
-        K_zeta_initial,
-        K_Z_array,
-        launch_angular_frequency,
-        mode_flag,
-        field.poloidal_flux,
-        find_density_1D,
-        field.B_R,
-        field.B_T,
-        field.B_Z,
-    )
-    H_other = find_H(
-        q_R_array,
-        q_Z_array,
-        K_R_array,
-        K_zeta_initial,
-        K_Z_array,
+    H_output = hamiltonian(q_R_array, q_Z_array, K_R_array, K_zeta_initial, K_Z_array)
+    # Create and immediately evaluate a Hamiltonian with the opposite mode
+    H_other = Hamiltonian(
+        field,
         launch_angular_frequency,
         -mode_flag,
-        field.poloidal_flux,
         find_density_1D,
-        field.B_R,
-        field.B_T,
-        field.B_Z,
-    )
-
-    # nabla H along the ray
-    dH_dR_output = find_dH_dR(
-        q_R_array,
-        q_Z_array,
-        K_R_array,
-        K_zeta_initial,
-        K_Z_array,
-        launch_angular_frequency,
-        mode_flag,
         delta_R,
-        field.poloidal_flux,
-        find_density_1D,
-        field.B_R,
-        field.B_T,
-        field.B_Z,
-    )
-    dH_dZ_output = find_dH_dZ(
-        q_R_array,
-        q_Z_array,
-        K_R_array,
-        K_zeta_initial,
-        K_Z_array,
-        launch_angular_frequency,
-        mode_flag,
         delta_Z,
-        field.poloidal_flux,
-        find_density_1D,
-        field.B_R,
-        field.B_T,
-        field.B_Z,
-    )
+        delta_K_R,
+        delta_K_zeta,
+        delta_K_Z,
+    )(q_R_array, q_Z_array, K_R_array, K_zeta_initial, K_Z_array)
 
     # Gradients of poloidal flux along the ray
     dpolflux_dR_debugging = find_dpolflux_dR(
@@ -1494,15 +991,6 @@ def event_reach_cutoff(
     dpolflux_dZ_debugging = find_dpolflux_dZ(
         q_R_array, q_Z_array, delta_Z, field.poloidal_flux
     )
-    # d2polflux_dR2_FFD_debugging =
-    # d2polflux_dZ2_FFD_debugging =
-
-    # Gradients of the total magnetic field along the ray
-    # dB_dR_FFD_debugging     =
-    # dB_dZ_FFD_debugging     =
-    # d2B_dR2_FFD_debugging   =
-    # d2B_dZ2_FFD_debugging   =
-    # d2B_dR_dZ_FFD_debugging =
 
     # -------------------
     # This saves the data generated by the main loop and the input data
@@ -2402,61 +1890,7 @@ def event_reach_cutoff(
     # print('ST first 2 terms / ST full: ', abs((G_term2[theta_m_min_idx]+G_term1[theta_m_min_idx])/G_full[theta_m_min_idx]) )
 
     # Calculates nabla nabla H, nabla_K nabla H, nabla_K nabla_K H
-    grad_grad_H = find_grad_grad_H_vectorised(
-        q_R_array,
-        q_Z_array,
-        K_R_array,
-        K_zeta_initial,
-        K_Z_array,
-        launch_angular_frequency,
-        mode_flag,
-        delta_R,
-        delta_Z,
-        field.poloidal_flux,
-        find_density_1D,
-        field.B_R,
-        field.B_T,
-        field.B_Z,
-    )
-
-    gradK_grad_H = find_gradK_grad_H_vectorised(
-        q_R_array,
-        q_Z_array,
-        K_R_array,
-        K_zeta_initial,
-        K_Z_array,
-        launch_angular_frequency,
-        mode_flag,
-        delta_K_R,
-        delta_K_zeta,
-        delta_K_Z,
-        delta_R,
-        delta_Z,
-        field.poloidal_flux,
-        find_density_1D,
-        field.B_R,
-        field.B_T,
-        field.B_Z,
-    )
-    grad_gradK_H = np.swapaxes(gradK_grad_H, 2, 1)
-
-    gradK_gradK_H = find_gradK_gradK_H_vectorised(
-        q_R_array,
-        q_Z_array,
-        K_R_array,
-        K_zeta_initial,
-        K_Z_array,
-        launch_angular_frequency,
-        mode_flag,
-        delta_K_R,
-        delta_K_zeta,
-        delta_K_Z,
-        field.poloidal_flux,
-        find_density_1D,
-        field.B_R,
-        field.B_T,
-        field.B_Z,
-    )
+    grad_grad_H, gradK_grad_H, gradK_gradK_H = hessians(dH)
 
     # gradK_gradK_H[:,0,1] = - gradK_gradK_H[:,0,1]
 
diff --git a/scotty/fun_CFD.py b/scotty/fun_CFD.py
index 08a150b3..e3a454c7 100644
--- a/scotty/fun_CFD.py
+++ b/scotty/fun_CFD.py
@@ -11,1390 +11,90 @@
 Run in Python 3,  does not work in Python 2
 """
 
-# import numpy as np
-# from scipy import constants as constants
-from scotty.fun_general import find_H
+import itertools
+from typing import Callable, Dict, Union, Tuple
 
+import numpy as np
 
-def find_dH_dR(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_R,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    H_plus = find_H(
-        q_R + delta_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_minus = find_H(
-        q_R - delta_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    dH_dR = (H_plus - H_minus) / (2 * delta_R)
-    return dH_dR
 
+WEIGHTS_AND_STEPS: Dict[int, Dict[str, list]] = {
+    1: {"weights": [-0.5, 0.5], "steps": [-1, 1]},
+    2: {"weights": [1, -2, 1], "steps": [-1, 0, 1]},
+}
 
-def find_dH_dZ(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_Z,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    H_plus = find_H(
-        q_R,
-        q_Z + delta_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_minus = find_H(
-        q_R,
-        q_Z - delta_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    dH_dZ = (H_plus - H_minus) / (2 * delta_Z)
-
-    return dH_dZ
 
-
-def find_dH_dKR(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_K_R,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
+def cfd_gradient(
+    f: Callable,
+    directions: Union[str, Tuple[str]],
+    dx: Union[float, Tuple[float]] = 1e-4,
+    d_order=1,
+    **kwargs,
 ):
-    H_plus = find_H(
-        q_R,
-        q_Z,
-        K_R + delta_K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_minus = find_H(
-        q_R,
-        q_Z,
-        K_R - delta_K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    dH_dKR = (H_plus - H_minus) / (2 * delta_K_R)
-
-    return dH_dKR
-
-
-def find_dH_dKzeta(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_K_zeta,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    H_plus = find_H(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta + delta_K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_minus = find_H(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta - delta_K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    dH_dKzeta = (H_plus - H_minus) / (2 * delta_K_zeta)
-
-    return dH_dKzeta
+    """Compute the derivative of a function using central finite differences
 
+    Arguments
+    ---------
+    f:
+        Function to compute derivative of
+    directions:
+        Tuple of directions to take derivative in. These must match
+        the name of the corresponding arguments to ``f``
+    dx:
+        Tuple of step sizes for each direction
+    d_order:
+        Order of the derivative (NOT the error scaling order of the
+        stencil)
+    **kwargs:
+        Arguments passed to ``f``
 
-def find_dH_dKZ(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_K_Z,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    H_plus = find_H(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z + delta_K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_minus = find_H(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z - delta_K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    dH_dKZ = (H_plus - H_minus) / (2 * delta_K_Z)
+    """
 
-    return dH_dKZ
+    if isinstance(directions, str):
+        directions = (directions,)
+    if isinstance(dx, float):
+        dx = (dx,)
 
+    num_directions = len(directions)
 
-def find_d2H_dR2(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_R,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    H_plus = find_H(
-        q_R + delta_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_current = find_H(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_minus = find_H(
-        q_R - delta_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dR2 = (H_plus - 2 * H_current + H_minus) / (delta_R**2)
+    if num_directions != len(dx):
+        raise ValueError(
+            f"Length of 'dx' ({len(dx)}) must match length of directions ({num_directions})"
+        )
+    if d_order < num_directions:
+        raise ValueError(
+            f"Order of derivative (d_order={d_order}) must be greater or equal to the number "
+            f"of directions ({num_directions})"
+        )
 
-    return d2H_dR2
+    d_order = d_order - (num_directions - 1)
 
+    try:
+        weights = WEIGHTS_AND_STEPS[d_order]["weights"]
+        steps = WEIGHTS_AND_STEPS[d_order]["steps"]
+    except KeyError:
+        raise ValueError(f"No weights for derivative order {d_order}")
 
-def find_d2H_dZ2(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_Z,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    H_plus = find_H(
-        q_R,
-        q_Z + delta_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_current = find_H(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_minus = find_H(
-        q_R,
-        q_Z - delta_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dZ2 = (H_plus - 2 * H_current + H_minus) / (delta_Z**2)
+    all_weights = itertools.product(*(weights,) * num_directions)
+    all_steps = itertools.product(*(steps,) * num_directions)
 
-    return d2H_dZ2
+    original_directions = {direction: kwargs[direction] for direction in directions}
 
+    result = 0.0
+    for weight, step in zip(all_weights, all_steps):
+        for direction, dx_, step_ in zip(directions, dx, step):
+            kwargs[direction] = original_directions[direction] + step_ * dx_
+        result += np.prod(weight) * f(**kwargs)
 
-def find_d2H_dR_dZ(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_R,
-    delta_Z,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    H_plus_R_plus_Z = find_H(
-        q_R + delta_R,
-        q_Z + delta_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_plus_R_minus_Z = find_H(
-        q_R + delta_R,
-        q_Z - delta_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_minus_R_plus_Z = find_H(
-        q_R - delta_R,
-        q_Z + delta_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_minus_R_minus_Z = find_H(
-        q_R - delta_R,
-        q_Z - delta_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dR_dZ = (
-        H_plus_R_plus_Z - H_plus_R_minus_Z - H_minus_R_plus_Z + H_minus_R_minus_Z
-    ) / (4 * delta_R * delta_Z)
-
-    return d2H_dR_dZ
-
-
-def find_d2H_dKR2(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_K_R,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    H_plus = find_H(
-        q_R,
-        q_Z,
-        K_R + delta_K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_current = find_H(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_minus = find_H(
-        q_R,
-        q_Z,
-        K_R - delta_K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKR2 = (H_plus - 2 * H_current + H_minus) / (delta_K_R**2)
-
-    return d2H_dKR2
-
-
-def find_d2H_dKR_dKzeta(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_K_R,
-    delta_K_zeta,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    H_plus_K_R_plus_K_zeta = find_H(
-        q_R,
-        q_Z,
-        K_R + delta_K_R,
-        K_zeta + delta_K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_plus_K_R_minus_K_zeta = find_H(
-        q_R,
-        q_Z,
-        K_R + delta_K_R,
-        K_zeta - delta_K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_minus_K_R_plus_K_zeta = find_H(
-        q_R,
-        q_Z,
-        K_R - delta_K_R,
-        K_zeta + delta_K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_minus_K_R_minus_K_zeta = find_H(
-        q_R,
-        q_Z,
-        K_R - delta_K_R,
-        K_zeta - delta_K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKR_dKzeta = (
-        H_plus_K_R_plus_K_zeta
-        - H_plus_K_R_minus_K_zeta
-        - H_minus_K_R_plus_K_zeta
-        + H_minus_K_R_minus_K_zeta
-    ) / (4 * delta_K_R * delta_K_zeta)
-
-    return d2H_dKR_dKzeta
-
-
-def find_d2H_dKR_dKZ(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_K_R,
-    delta_K_Z,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    H_plus_K_R_plus_K_z = find_H(
-        q_R,
-        q_Z,
-        K_R + delta_K_R,
-        K_zeta,
-        K_Z + delta_K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_plus_K_R_minus_K_z = find_H(
-        q_R,
-        q_Z,
-        K_R + delta_K_R,
-        K_zeta,
-        K_Z - delta_K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_minus_K_R_plus_K_z = find_H(
-        q_R,
-        q_Z,
-        K_R - delta_K_R,
-        K_zeta,
-        K_Z + delta_K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_minus_K_R_minus_K_z = find_H(
-        q_R,
-        q_Z,
-        K_R - delta_K_R,
-        K_zeta,
-        K_Z - delta_K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKR_dKz = (
-        H_plus_K_R_plus_K_z
-        - H_plus_K_R_minus_K_z
-        - H_minus_K_R_plus_K_z
-        + H_minus_K_R_minus_K_z
-    ) / (4 * delta_K_R * delta_K_Z)
-
-    return d2H_dKR_dKz
-
-
-def find_d2H_dKzeta2(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_K_zeta,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    H_plus = find_H(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta + delta_K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_current = find_H(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_minus = find_H(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta - delta_K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKzeta2 = (H_plus - 2 * H_current + H_minus) / (delta_K_zeta**2)
-
-    return d2H_dKzeta2
-
-
-def find_d2H_dKzeta_dKZ(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_K_zeta,
-    delta_K_Z,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    H_plus_K_zeta_plus_K_z = find_H(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta + delta_K_zeta,
-        K_Z + delta_K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_plus_K_zeta_minus_K_z = find_H(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta + delta_K_zeta,
-        K_Z - delta_K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_minus_K_zeta_plus_K_z = find_H(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta - delta_K_zeta,
-        K_Z + delta_K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_minus_K_zeta_minus_K_z = find_H(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta - delta_K_zeta,
-        K_Z - delta_K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKzeta_dKz = (
-        H_plus_K_zeta_plus_K_z
-        - H_plus_K_zeta_minus_K_z
-        - H_minus_K_zeta_plus_K_z
-        + H_minus_K_zeta_minus_K_z
-    ) / (4 * delta_K_zeta * delta_K_Z)
-
-    return d2H_dKzeta_dKz
-
-
-def find_d2H_dKZ2(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_K_Z,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    H_plus = find_H(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z + delta_K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_current = find_H(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_minus = find_H(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z - delta_K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKZ2 = (H_plus - 2 * H_current + H_minus) / (delta_K_Z**2)
-
-    return d2H_dKZ2
-
-
-def find_d2H_dKR_dR(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_K_R,
-    delta_R,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    H_plus_K_R_plus_R = find_H(
-        q_R + delta_R,
-        q_Z,
-        K_R + delta_K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_plus_K_R_minus_R = find_H(
-        q_R - delta_R,
-        q_Z,
-        K_R + delta_K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_minus_K_R_plus_R = find_H(
-        q_R + delta_R,
-        q_Z,
-        K_R - delta_K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_minus_K_R_minus_R = find_H(
-        q_R - delta_R,
-        q_Z,
-        K_R - delta_K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKR_dR = (
-        H_plus_K_R_plus_R
-        - H_plus_K_R_minus_R
-        - H_minus_K_R_plus_R
-        + H_minus_K_R_minus_R
-    ) / (4 * delta_K_R * delta_R)
-
-    return d2H_dKR_dR
-
-
-def find_d2H_dKR_dZ(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_K_R,
-    delta_Z,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    H_plus_K_R_plus_Z = find_H(
-        q_R,
-        q_Z + delta_Z,
-        K_R + delta_K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_plus_K_R_minus_Z = find_H(
-        q_R,
-        q_Z - delta_Z,
-        K_R + delta_K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_minus_K_R_plus_Z = find_H(
-        q_R,
-        q_Z + delta_Z,
-        K_R - delta_K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_minus_K_R_minus_Z = find_H(
-        q_R,
-        q_Z - delta_Z,
-        K_R - delta_K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKR_dZ = (
-        H_plus_K_R_plus_Z
-        - H_plus_K_R_minus_Z
-        - H_minus_K_R_plus_Z
-        + H_minus_K_R_minus_Z
-    ) / (4 * delta_K_R * delta_Z)
-
-    return d2H_dKR_dZ
-
-
-def find_d2H_dKzeta_dR(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_K_zeta,
-    delta_R,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    H_plus_K_zeta_plus_R = find_H(
-        q_R + delta_R,
-        q_Z,
-        K_R,
-        K_zeta + delta_K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_plus_K_zeta_minus_R = find_H(
-        q_R - delta_R,
-        q_Z,
-        K_R,
-        K_zeta + delta_K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_minus_K_zeta_plus_R = find_H(
-        q_R + delta_R,
-        q_Z,
-        K_R,
-        K_zeta - delta_K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_minus_K_zeta_minus_R = find_H(
-        q_R - delta_R,
-        q_Z,
-        K_R,
-        K_zeta - delta_K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKzeta_dR = (
-        H_plus_K_zeta_plus_R
-        - H_plus_K_zeta_minus_R
-        - H_minus_K_zeta_plus_R
-        + H_minus_K_zeta_minus_R
-    ) / (4 * delta_K_zeta * delta_R)
-
-    return d2H_dKzeta_dR
-
-
-def find_d2H_dKzeta_dZ(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_K_zeta,
-    delta_Z,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    H_plus_K_zeta_plus_Z = find_H(
-        q_R,
-        q_Z + delta_Z,
-        K_R,
-        K_zeta + delta_K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_plus_K_zeta_minus_Z = find_H(
-        q_R,
-        q_Z - delta_Z,
-        K_R,
-        K_zeta + delta_K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_minus_K_zeta_plus_Z = find_H(
-        q_R,
-        q_Z + delta_Z,
-        K_R,
-        K_zeta - delta_K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_minus_K_zeta_minus_Z = find_H(
-        q_R,
-        q_Z - delta_Z,
-        K_R,
-        K_zeta - delta_K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKzeta_dZ = (
-        H_plus_K_zeta_plus_Z
-        - H_plus_K_zeta_minus_Z
-        - H_minus_K_zeta_plus_Z
-        + H_minus_K_zeta_minus_Z
-    ) / (4 * delta_K_zeta * delta_Z)
-
-    return d2H_dKzeta_dZ
-
-
-def find_d2H_dKZ_dR(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_K_Z,
-    delta_R,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    H_plus_K_Z_plus_R = find_H(
-        q_R + delta_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z + delta_K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_plus_K_Z_minus_R = find_H(
-        q_R - delta_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z + delta_K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_minus_K_Z_plus_R = find_H(
-        q_R + delta_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z - delta_K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_minus_K_Z_minus_R = find_H(
-        q_R - delta_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z - delta_K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKZ_dR = (
-        H_plus_K_Z_plus_R
-        - H_plus_K_Z_minus_R
-        - H_minus_K_Z_plus_R
-        + H_minus_K_Z_minus_R
-    ) / (4 * delta_K_Z * delta_R)
-
-    return d2H_dKZ_dR
-
-
-def find_d2H_dKZ_dZ(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_K_Z,
-    delta_Z,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    H_plus_K_Z_plus_Z = find_H(
-        q_R,
-        q_Z + delta_Z,
-        K_R,
-        K_zeta,
-        K_Z + delta_K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_plus_K_Z_minus_Z = find_H(
-        q_R,
-        q_Z - delta_Z,
-        K_R,
-        K_zeta,
-        K_Z + delta_K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_minus_K_Z_plus_Z = find_H(
-        q_R,
-        q_Z + delta_Z,
-        K_R,
-        K_zeta,
-        K_Z - delta_K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_minus_K_Z_minus_Z = find_H(
-        q_R,
-        q_Z - delta_Z,
-        K_R,
-        K_zeta,
-        K_Z - delta_K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKZ_dZ = (
-        H_plus_K_Z_plus_Z
-        - H_plus_K_Z_minus_Z
-        - H_minus_K_Z_plus_Z
-        + H_minus_K_Z_minus_Z
-    ) / (4 * delta_K_Z * delta_Z)
-
-    return d2H_dKZ_dZ
+    return result / np.prod(dx * d_order)
 
 
 def find_dpolflux_dR(q_R, q_Z, delta_R, interp_poloidal_flux):
-    polflux_plus = interp_poloidal_flux(q_R + delta_R, q_Z)
-    polflux_minus = interp_poloidal_flux(q_R - delta_R, q_Z)
-
-    dpolflux_dR = (polflux_plus - polflux_minus) / (2 * delta_R)
-
-    return dpolflux_dR
+    return cfd_gradient(
+        interp_poloidal_flux, q_R=q_R, q_Z=q_Z, directions="q_R", dx=delta_R
+    )
 
 
 def find_dpolflux_dZ(q_R, q_Z, delta_Z, interp_poloidal_flux):
-    polflux_plus = interp_poloidal_flux(q_R, q_Z + delta_Z)
-    polflux_minus = interp_poloidal_flux(q_R, q_Z - delta_Z)
-
-    dpolflux_dZ = (polflux_plus - polflux_minus) / (2 * delta_Z)
-
-    return dpolflux_dZ
+    return cfd_gradient(
+        interp_poloidal_flux, q_R=q_R, q_Z=q_Z, directions="q_Z", dx=delta_Z
+    )
diff --git a/scotty/fun_FFD.py b/scotty/fun_FFD.py
deleted file mode 100644
index 9f6e9be0..00000000
--- a/scotty/fun_FFD.py
+++ /dev/null
@@ -1,1406 +0,0 @@
-# -*- coding: utf-8 -*-
-"""
-Created on Fri Jun  8 10:44:34 2018
-
-Functions for finding derivatives of H using forward finite difference.
-
-@author: chenv
-Valerian Hongjie Hall-Chen
-valerian_hall-chen@ihpc.a-star.edu.sg
-
-Run in Python 3,  does not work in Python 2
-"""
-
-# import numpy as np
-# from scipy import constants as constants
-from scotty.fun_general import find_H
-
-
-def find_dH_dR(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_R,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    H_0 = find_H(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_1 = find_H(
-        q_R + delta_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_2 = find_H(
-        q_R + 2 * delta_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    dH_dR = ((-3 / 2) * H_0 + (2) * H_1 + (-1 / 2) * H_2) / (delta_R)
-
-    return dH_dR
-
-
-def find_dH_dZ(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_Z,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    H_0 = find_H(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_1 = find_H(
-        q_R,
-        q_Z + delta_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_2 = find_H(
-        q_R,
-        q_Z + 2 * delta_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    dH_dZ = ((-3 / 2) * H_0 + (2) * H_1 + (-1 / 2) * H_2) / (delta_Z)
-
-    return dH_dZ
-
-
-def find_dH_dKR(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_K_R,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    H_0 = find_H(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_1 = find_H(
-        q_R,
-        q_Z,
-        K_R + delta_K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_2 = find_H(
-        q_R,
-        q_Z,
-        K_R + 2 * delta_K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    dH_dKR = ((-3 / 2) * H_0 + (2) * H_1 + (-1 / 2) * H_2) / (delta_K_R)
-
-    return dH_dKR
-
-
-def find_dH_dKzeta(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_K_zeta,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    H_0 = find_H(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_1 = find_H(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta + delta_K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_2 = find_H(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta + 2 * delta_K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    dH_dKzeta = ((-3 / 2) * H_0 + (2) * H_1 + (-1 / 2) * H_2) / (delta_K_zeta)
-
-    return dH_dKzeta
-
-
-def find_dH_dKZ(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_K_Z,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    H_0 = find_H(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_1 = find_H(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z + delta_K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_2 = find_H(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z + 2 * delta_K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    dH_dKZ = ((-3 / 2) * H_0 + (2) * H_1 + (-1 / 2) * H_2) / (delta_K_Z)
-
-    return dH_dKZ
-
-
-def find_d2H_dR2(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_R,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    H_0 = find_H(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_1 = find_H(
-        q_R + delta_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_2 = find_H(
-        q_R + 2 * delta_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_3 = find_H(
-        q_R + 3 * delta_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    #    H_4 = find_H(q_R+4*delta_R,
-    #                 q_Z, K_R, K_zeta, K_Z, launch_angular_frequency, mode_flag,
-    #                 interp_poloidal_flux, find_density_1D, find_B_R, find_B_T, find_B_Z)
-    d2H_dR2 = ((2) * H_0 + (-5) * H_1 + (4) * H_2 + (-1) * H_3) / (delta_R**2)
-    #    d2H_dR2 = ( (1)*H_0 + (-2)*H_1 + (1)*H_2 ) / (delta_R**2)
-
-    return d2H_dR2
-
-
-def find_d2H_dZ2(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_Z,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    H_0 = find_H(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_1 = find_H(
-        q_R,
-        q_Z + delta_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_2 = find_H(
-        q_R,
-        q_Z + 2 * delta_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_3 = find_H(
-        q_R,
-        q_Z + 3 * delta_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dZ2 = ((2) * H_0 + (-5) * H_1 + (4) * H_2 + (-1) * H_3) / (delta_Z**2)
-
-    return d2H_dZ2
-
-
-def find_d2H_dR_dZ(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_R,
-    delta_Z,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    dH_dZ_0 = find_dH_dZ(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    dH_dZ_1 = find_dH_dZ(
-        q_R + delta_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    dH_dZ_2 = find_dH_dZ(
-        q_R + 2 * delta_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dR_dZ = ((-3 / 2) * dH_dZ_0 + (2) * dH_dZ_1 + (-1 / 2) * dH_dZ_2) / (delta_R)
-
-    return d2H_dR_dZ
-
-
-def find_d2H_dKR2(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_K_R,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    H_0 = find_H(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_1 = find_H(
-        q_R,
-        q_Z,
-        K_R + delta_K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_2 = find_H(
-        q_R,
-        q_Z,
-        K_R + 2 * delta_K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_3 = find_H(
-        q_R,
-        q_Z,
-        K_R + 3 * delta_K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKR2 = ((2) * H_0 + (-5) * H_1 + (4) * H_2 + (-1) * H_3) / (delta_K_R**2)
-
-    return d2H_dKR2
-
-
-def find_d2H_dKR_dKzeta(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_K_R,
-    delta_K_zeta,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    dH_dKzeta_0 = find_dH_dKzeta(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_K_zeta,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    dH_dKzeta_1 = find_dH_dKzeta(
-        q_R,
-        q_Z,
-        K_R + delta_K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_K_zeta,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    dH_dKzeta_2 = find_dH_dKzeta(
-        q_R,
-        q_Z,
-        K_R + 2 * delta_K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_K_zeta,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKR_dKzeta = (
-        (-3 / 2) * dH_dKzeta_0 + (2) * dH_dKzeta_1 + (-1 / 2) * dH_dKzeta_2
-    ) / (delta_K_R)
-
-    return d2H_dKR_dKzeta
-
-
-def find_d2H_dKR_dKZ(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_K_R,
-    delta_K_Z,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    dH_dKZ_0 = find_dH_dKZ(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_K_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    dH_dKZ_1 = find_dH_dKZ(
-        q_R,
-        q_Z,
-        K_R + delta_K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_K_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    dH_dKZ_2 = find_dH_dKZ(
-        q_R,
-        q_Z,
-        K_R + 2 * delta_K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_K_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKR_dKz = ((-3 / 2) * dH_dKZ_0 + (2) * dH_dKZ_1 + (-1 / 2) * dH_dKZ_2) / (
-        delta_K_R
-    )
-
-    return d2H_dKR_dKz
-
-
-def find_d2H_dKzeta2(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_K_zeta,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    H_0 = find_H(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_1 = find_H(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta + delta_K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_2 = find_H(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta + 2 * delta_K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_3 = find_H(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta + 3 * delta_K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKzeta2 = ((2) * H_0 + (-5) * H_1 + (4) * H_2 + (-1) * H_3) / (
-        delta_K_zeta**2
-    )
-
-    return d2H_dKzeta2
-
-
-def find_d2H_dKzeta_dKZ(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_K_zeta,
-    delta_K_Z,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    dH_dKZ_0 = find_dH_dKZ(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_K_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    dH_dKZ_1 = find_dH_dKZ(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta + delta_K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_K_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    dH_dKZ_2 = find_dH_dKZ(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta + 2 * delta_K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_K_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKzeta_dKZ = ((-3 / 2) * dH_dKZ_0 + (2) * dH_dKZ_1 + (-1 / 2) * dH_dKZ_2) / (
-        delta_K_zeta
-    )
-
-    return d2H_dKzeta_dKZ
-
-
-def find_d2H_dKZ2(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_K_Z,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    H_0 = find_H(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_1 = find_H(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z + delta_K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_2 = find_H(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z + 2 * delta_K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    H_3 = find_H(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z + 3 * delta_K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKZ2 = ((2) * H_0 + (-5) * H_1 + (4) * H_2 + (-1) * H_3) / (delta_K_Z**2)
-
-    return d2H_dKZ2
-
-
-def find_d2H_dKR_dR(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_K_R,
-    delta_R,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    dH_dR_0 = find_dH_dR(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_R,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    dH_dR_1 = find_dH_dR(
-        q_R,
-        q_Z,
-        K_R + delta_K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_R,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    dH_dR_2 = find_dH_dR(
-        q_R,
-        q_Z,
-        K_R + 2 * delta_K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_R,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKR_dR = ((-3 / 2) * dH_dR_0 + (2) * dH_dR_1 + (-1 / 2) * dH_dR_2) / (delta_K_R)
-
-    return d2H_dKR_dR
-
-
-def find_d2H_dKR_dZ(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_K_R,
-    delta_Z,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    dH_dZ_0 = find_dH_dZ(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    dH_dZ_1 = find_dH_dZ(
-        q_R,
-        q_Z,
-        K_R + delta_K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    dH_dZ_2 = find_dH_dZ(
-        q_R,
-        q_Z,
-        K_R + 2 * delta_K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKR_dZ = ((-3 / 2) * dH_dZ_0 + (2) * dH_dZ_1 + (-1 / 2) * dH_dZ_2) / (delta_K_R)
-
-    return d2H_dKR_dZ
-
-
-def find_d2H_dKzeta_dR(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_K_zeta,
-    delta_R,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    dH_dR_0 = find_dH_dR(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_R,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    dH_dR_1 = find_dH_dR(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta + delta_K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_R,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    dH_dR_2 = find_dH_dR(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta + 2 * delta_K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_R,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKzeta_dR = ((-3 / 2) * dH_dR_0 + (2) * dH_dR_1 + (-1 / 2) * dH_dR_2) / (
-        delta_K_zeta
-    )
-
-    return d2H_dKzeta_dR
-
-
-def find_d2H_dKzeta_dZ(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_K_zeta,
-    delta_Z,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    dH_dZ_0 = find_dH_dZ(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    dH_dZ_1 = find_dH_dZ(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta + delta_K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    dH_dZ_2 = find_dH_dZ(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta + 2 * delta_K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKzeta_dZ = ((-3 / 2) * dH_dZ_0 + (2) * dH_dZ_1 + (-1 / 2) * dH_dZ_2) / (
-        delta_K_zeta
-    )
-
-    return d2H_dKzeta_dZ
-
-
-def find_d2H_dKZ_dR(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_K_Z,
-    delta_R,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    dH_dR_0 = find_dH_dR(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_R,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    dH_dR_1 = find_dH_dR(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z + delta_K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_R,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    dH_dR_2 = find_dH_dR(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z + 2 * delta_K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_R,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKZ_dR = ((-3 / 2) * dH_dR_0 + (2) * dH_dR_1 + (-1 / 2) * dH_dR_2) / (delta_K_Z)
-
-    return d2H_dKZ_dR
-
-
-def find_d2H_dKZ_dZ(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_K_Z,
-    delta_Z,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    dH_dZ_0 = find_dH_dZ(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    dH_dZ_1 = find_dH_dZ(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z + delta_K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    dH_dZ_2 = find_dH_dZ(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z + 2 * delta_K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKZ_dZ = ((-3 / 2) * dH_dZ_0 + (2) * dH_dZ_1 + (-1 / 2) * dH_dZ_2) / (delta_K_Z)
-
-    return d2H_dKZ_dZ
-
-
-def find_dpolflux_dR(q_R, q_Z, delta_R, interp_poloidal_flux):
-    polflux_0 = interp_poloidal_flux(q_R, q_Z)
-    polflux_1 = interp_poloidal_flux(q_R + delta_R, q_Z)
-    polflux_2 = interp_poloidal_flux(q_R + 2 * delta_R, q_Z)
-
-    dpolflux_dR = ((-3 / 2) * polflux_0 + (2) * polflux_1 + (-1 / 2) * polflux_2) / (
-        delta_R
-    )
-
-    return dpolflux_dR
-
-
-def find_dpolflux_dZ(q_R, q_Z, delta_Z, interp_poloidal_flux):
-    polflux_0 = interp_poloidal_flux(q_R, q_Z)
-    polflux_1 = interp_poloidal_flux(q_R, q_Z + delta_Z)
-    polflux_2 = interp_poloidal_flux(q_R, q_Z + 2 * delta_Z)
-
-    dpolflux_dZ = ((-3 / 2) * polflux_0 + (2) * polflux_1 + (-1 / 2) * polflux_2) / (
-        delta_Z
-    )
-
-    return dpolflux_dZ
diff --git a/scotty/fun_evolution.py b/scotty/fun_evolution.py
index af1cdc48..de4ed50f 100644
--- a/scotty/fun_evolution.py
+++ b/scotty/fun_evolution.py
@@ -14,398 +14,14 @@
 
 import numpy as np
 
-from scotty.fun_FFD import find_dH_dR, find_dH_dZ  # \nabla H
-from scotty.fun_CFD import find_dH_dKR, find_dH_dKZ, find_dH_dKzeta  # \nabla_K H
-from scotty.fun_FFD import find_d2H_dR2, find_d2H_dZ2, find_d2H_dR_dZ  # \nabla \nabla H
-from scotty.fun_CFD import (
-    find_d2H_dKR2,
-    find_d2H_dKR_dKzeta,
-    find_d2H_dKR_dKZ,
-    find_d2H_dKzeta2,
-    find_d2H_dKzeta_dKZ,
-    find_d2H_dKZ2,
-)  # \nabla_K \nabla_K H
-from scotty.fun_mix import (
-    find_d2H_dKR_dR,
-    find_d2H_dKR_dZ,
-    find_d2H_dKzeta_dR,
-    find_d2H_dKzeta_dZ,
-    find_d2H_dKZ_dR,
-    find_d2H_dKZ_dZ,
-)  # \nabla_K \nabla H
-
-from scotty.fun_general import find_H_numba
-
-
-# Functions (gradients of H, vectorised)
-def find_grad_grad_H_vectorised(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_R,
-    delta_Z,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    # \nabla \nabla H
-    d2H_dR2 = find_d2H_dR2(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_R,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dR_dZ = find_d2H_dR_dZ(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_R,
-        delta_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dZ2 = find_d2H_dZ2(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-
-    zeros = np.zeros_like(d2H_dR2)
-    grad_grad_H = np.moveaxis(
-        np.squeeze(
-            np.array(
-                [
-                    [d2H_dR2, zeros, d2H_dR_dZ],
-                    [zeros, zeros, zeros],
-                    [d2H_dR_dZ, zeros, d2H_dZ2],
-                ]
-            )
-        ),
-        2,
-        0,
-    )  # Such that shape is [points,3,3] instead of [3,3,points]
-
-    return grad_grad_H
-
-
-def find_gradK_grad_H_vectorised(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_K_R,
-    delta_K_zeta,
-    delta_K_Z,
-    delta_R,
-    delta_Z,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    # \nabla_K \nabla H
-    d2H_dKR_dR = find_d2H_dKR_dR(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_K_R,
-        delta_R,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKZ_dZ = find_d2H_dKZ_dZ(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_K_Z,
-        delta_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKR_dZ = find_d2H_dKR_dZ(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_K_R,
-        delta_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKzeta_dZ = find_d2H_dKzeta_dZ(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_K_zeta,
-        delta_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKzeta_dR = find_d2H_dKzeta_dR(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_K_zeta,
-        delta_R,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKZ_dR = find_d2H_dKZ_dR(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_K_Z,
-        delta_R,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-
-    zeros = np.zeros_like(d2H_dKR_dR)
-    gradK_grad_H = np.moveaxis(
-        np.squeeze(
-            np.array(
-                [
-                    [d2H_dKR_dR, zeros, d2H_dKR_dZ],
-                    [d2H_dKzeta_dR, zeros, d2H_dKzeta_dZ],
-                    [d2H_dKZ_dR, zeros, d2H_dKZ_dZ],
-                ]
-            )
-        ),
-        2,
-        0,
-    )  # Such that shape is [points,3,3] instead of [3,3,points]
-
-    return gradK_grad_H
-
-
-def find_gradK_gradK_H_vectorised(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_K_R,
-    delta_K_zeta,
-    delta_K_Z,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    # \nabla_K \nabla_K H
-    d2H_dKR2 = find_d2H_dKR2(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_K_R,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKzeta2 = find_d2H_dKzeta2(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_K_zeta,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKZ2 = find_d2H_dKZ2(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_K_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKR_dKzeta = find_d2H_dKR_dKzeta(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_K_R,
-        delta_K_zeta,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKR_dKZ = find_d2H_dKR_dKZ(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_K_R,
-        delta_K_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKzeta_dKZ = find_d2H_dKzeta_dKZ(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_K_zeta,
-        delta_K_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    gradK_gradK_H = np.moveaxis(
-        np.squeeze(
-            np.array(
-                [
-                    [d2H_dKR2, d2H_dKR_dKzeta, d2H_dKR_dKZ],
-                    [d2H_dKR_dKzeta, d2H_dKzeta2, d2H_dKzeta_dKZ],
-                    [d2H_dKR_dKZ, d2H_dKzeta_dKZ, d2H_dKZ2],
-                ]
-            )
-        ),
-        2,
-        0,
-    )  # Such that shape is [points,3,3] instead of [3,3,points]
-
-    return gradK_gradK_H
-
-
-# ---------------------------------
+from scotty.hamiltonian import Hamiltonian, hessians
+from scotty.typing import FloatArray
 
 
 # Functions (solver)
 # Defines the ray evolution function
 # Not necessary for what I want to do, but it does make life easier
-def ray_evolution_2D_fun(
-    tau,
-    ray_parameters_2D,
-    K_zeta,
-    launch_angular_frequency,
-    mode_flag,
-    delta_R,
-    delta_Z,
-    delta_K_R,
-    delta_K_zeta,
-    delta_K_Z,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
+def ray_evolution_2D_fun(tau, ray_parameters_2D, K_zeta, hamiltonian: Hamiltonian):
     """
     Parameters
     ----------
@@ -413,6 +29,8 @@ def ray_evolution_2D_fun(
         Parameter along the ray.
     ray_parameters_2D : complex128
         q_R, q_Z, K_R, K_Z
+    hamiltonian:
+        Hamiltonian object
 
     Returns
     -------
@@ -424,108 +42,89 @@ def ray_evolution_2D_fun(
 
     """
 
-    ## Clean input up. Not necessary, but aids readability
+    # Clean input up. Not necessary, but aids readability
     q_R = ray_parameters_2D[0]
     q_Z = ray_parameters_2D[1]
     K_R = ray_parameters_2D[2]
     K_Z = ray_parameters_2D[3]
 
-    ## Find derivatives of H
-    # \nabla H
-    dH_dR = find_dH_dR(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_R,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    dH_dZ = find_dH_dZ(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-
-    # \nabla_K H
-    dH_dKR = find_dH_dKR(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_K_R,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    dH_dKZ = find_dH_dKZ(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_K_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
+    # Find derivatives of H
+    dH = hamiltonian.derivatives(q_R, q_Z, K_R, K_zeta, K_Z)
 
     d_ray_parameters_2D_d_tau = np.zeros_like(ray_parameters_2D)
 
-    d_ray_parameters_2D_d_tau[0] = dH_dKR  # d (q_R) / d tau
-    d_ray_parameters_2D_d_tau[1] = dH_dKZ  # d (q_Z) / d tau
-    d_ray_parameters_2D_d_tau[2] = -dH_dR  # d (K_R) / d tau
-    d_ray_parameters_2D_d_tau[3] = -dH_dZ  # d (K_Z) / d tau
+    # d (q_R) / d tau
+    d_ray_parameters_2D_d_tau[0] = dH["dH_dKR"]
+    # d (q_Z) / d tau
+    d_ray_parameters_2D_d_tau[1] = dH["dH_dKZ"]
+    # d (K_R) / d tau
+    d_ray_parameters_2D_d_tau[2] = -dH["dH_dR"]
+    # d (K_Z) / d tau
+    d_ray_parameters_2D_d_tau[3] = -dH["dH_dZ"]
 
     return d_ray_parameters_2D_d_tau
 
 
-# -------------------
-
-
-# Defines the beam evolution function
-def beam_evolution_fun(
-    tau,
-    beam_parameters,
-    K_zeta,
-    launch_angular_frequency,
-    mode_flag,
-    delta_R,
-    delta_Z,
-    delta_K_R,
-    delta_K_zeta,
-    delta_K_Z,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
+def pack_beam_parameters(q_R, q_zeta, q_Z, K_R, K_Z, Psi):
+    # This used to be complex, with a length of 11, but the solver
+    # throws a warning saying that something is casted to real It
+    # seems to be fine, bu
+    beam_parameters = np.zeros(17)
+
+    beam_parameters[0] = q_R
+    beam_parameters[1] = q_zeta
+    beam_parameters[2] = q_Z
+
+    beam_parameters[3] = K_R
+    beam_parameters[4] = K_Z
+
+    beam_parameters[5] = np.real(Psi[0, 0])
+    beam_parameters[6] = np.real(Psi[1, 1])
+    beam_parameters[7] = np.real(Psi[2, 2])
+    beam_parameters[8] = np.real(Psi[0, 1])
+    beam_parameters[9] = np.real(Psi[0, 2])
+    beam_parameters[10] = np.real(Psi[1, 2])
+
+    beam_parameters[11] = np.imag(Psi[0, 0])
+    beam_parameters[12] = np.imag(Psi[1, 1])
+    beam_parameters[13] = np.imag(Psi[2, 2])
+    beam_parameters[14] = np.imag(Psi[0, 1])
+    beam_parameters[15] = np.imag(Psi[0, 2])
+    beam_parameters[16] = np.imag(Psi[1, 2])
+    return beam_parameters
+
+
+def unpack_beam_parameters(beam_parameters: FloatArray):
+    q_R = beam_parameters[0, ...]
+    q_zeta = beam_parameters[1, ...]
+    q_Z = beam_parameters[2, ...]
+    K_R = beam_parameters[3, ...]
+    K_Z = beam_parameters[4, ...]
+
+    n_points = beam_parameters.shape[1] if beam_parameters.ndim == 2 else 1
+
+    Psi_3D = np.zeros([n_points, 3, 3], dtype="complex128")
+    # Psi_RR
+    Psi_3D[:, 0, 0] = beam_parameters[5, ...] + 1j * beam_parameters[11, ...]
+    # Psi_zetazeta
+    Psi_3D[:, 1, 1] = beam_parameters[6, ...] + 1j * beam_parameters[12, ...]
+    # Psi_ZZ
+    Psi_3D[:, 2, 2] = beam_parameters[7, ...] + 1j * beam_parameters[13, ...]
+    # Psi_Rzeta
+    Psi_3D[:, 0, 1] = beam_parameters[8, ...] + 1j * beam_parameters[14, ...]
+    # Psi_RZ
+    Psi_3D[:, 0, 2] = beam_parameters[9, ...] + 1j * beam_parameters[15, ...]
+    # Psi_zetaZ
+    Psi_3D[:, 1, 2] = beam_parameters[10, ...] + 1j * beam_parameters[16, ...]
+    # Psi_3D is symmetric
+    Psi_3D[:, 1, 0] = Psi_3D[:, 0, 1]
+    Psi_3D[:, 2, 0] = Psi_3D[:, 0, 2]
+    Psi_3D[:, 2, 1] = Psi_3D[:, 1, 2]
+
+    return q_R, q_zeta, q_Z, K_R, K_Z, np.squeeze(Psi_3D)
+
+
+def beam_evolution_fun(tau, beam_parameters, K_zeta, hamiltonian: Hamiltonian):
     """
     Parameters
     ----------
@@ -540,376 +139,19 @@ def beam_evolution_fun(
         d (beam_parameters) / d tau
     """
 
-    ## Clean input up. Not necessary, but aids readability
-    q_R = beam_parameters[0]
-    q_zeta = beam_parameters[1]
-    q_Z = beam_parameters[2]
-    K_R = beam_parameters[3]
-    K_Z = beam_parameters[4]
-
-    Psi_3D = np.zeros([3, 3], dtype="complex128")
-    Psi_3D[0, 0] = beam_parameters[5] + 1j * beam_parameters[11]  # Psi_RR
-    Psi_3D[1, 1] = beam_parameters[6] + 1j * beam_parameters[12]  # Psi_zetazeta
-    Psi_3D[2, 2] = beam_parameters[7] + 1j * beam_parameters[13]  # Psi_ZZ
-    Psi_3D[0, 1] = beam_parameters[8] + 1j * beam_parameters[14]  # Psi_Rzeta
-    Psi_3D[0, 2] = beam_parameters[9] + 1j * beam_parameters[15]  # Psi_RZ
-    Psi_3D[1, 2] = beam_parameters[10] + 1j * beam_parameters[16]  # Psi_zetaZ
-    Psi_3D[1, 0] = Psi_3D[0, 1]  # Psi_3D is symmetric
-    Psi_3D[2, 0] = Psi_3D[0, 2]
-    Psi_3D[2, 1] = Psi_3D[1, 2]
-
-    ## Find derivatives of H
-    # \nabla H
-    dH_dR = find_dH_dR(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_R,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    dH_dZ = find_dH_dZ(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-
-    # \nabla_K H
-    dH_dKR = find_dH_dKR(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_K_R,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    dH_dKzeta = find_dH_dKzeta(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_K_zeta,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    dH_dKZ = find_dH_dKZ(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_K_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
+    q_R, q_zeta, q_Z, K_R, K_Z, Psi_3D = unpack_beam_parameters(beam_parameters)
 
-    # \nabla \nabla H
-    d2H_dR2 = find_d2H_dR2(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_R,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dR_dZ = find_d2H_dR_dZ(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_R,
-        delta_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dZ2 = find_d2H_dZ2(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    grad_grad_H = np.squeeze(
-        np.array(
-            [
-                [d2H_dR2.item(), 0.0, d2H_dR_dZ.item()],
-                [0.0, 0.0, 0.0],
-                [
-                    d2H_dR_dZ.item(),
-                    0.0,
-                    d2H_dZ2.item(),
-                ],  # . item() to convert variable from type ndarray to float, such that the array elements all have the same type
-            ]
-        )
-    )
+    # Find derivatives of H
+    dH = hamiltonian.derivatives(q_R, q_Z, K_R, K_zeta, K_Z, second_order=True)
 
-    # \nabla_K \nabla H
-    d2H_dKR_dR = find_d2H_dKR_dR(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_K_R,
-        delta_R,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKZ_dZ = find_d2H_dKZ_dZ(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_K_Z,
-        delta_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKR_dZ = find_d2H_dKR_dZ(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_K_R,
-        delta_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKzeta_dZ = find_d2H_dKzeta_dZ(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_K_zeta,
-        delta_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKzeta_dR = find_d2H_dKzeta_dR(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_K_zeta,
-        delta_R,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKZ_dR = find_d2H_dKZ_dR(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_K_Z,
-        delta_R,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    gradK_grad_H = np.squeeze(
-        np.array(
-            [
-                [d2H_dKR_dR.item(), 0.0, d2H_dKR_dZ.item()],
-                [d2H_dKzeta_dR.item(), 0.0, d2H_dKzeta_dZ.item()],
-                [d2H_dKZ_dR.item(), 0.0, d2H_dKZ_dZ.item()],
-            ]
-        )
-    )
+    grad_grad_H, gradK_grad_H, gradK_gradK_H = hessians(dH)
     grad_gradK_H = np.transpose(gradK_grad_H)
 
-    # \nabla_K \nabla_K H
-    d2H_dKR2 = find_d2H_dKR2(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_K_R,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKzeta2 = find_d2H_dKzeta2(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_K_zeta,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKZ2 = find_d2H_dKZ2(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_K_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKR_dKzeta = find_d2H_dKR_dKzeta(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_K_R,
-        delta_K_zeta,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKR_dKZ = find_d2H_dKR_dKZ(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_K_R,
-        delta_K_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKzeta_dKZ = find_d2H_dKzeta_dKZ(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_K_zeta,
-        delta_K_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    gradK_gradK_H = np.squeeze(
-        np.array(
-            [
-                [d2H_dKR2.item(), d2H_dKR_dKzeta.item(), d2H_dKR_dKZ.item()],
-                [d2H_dKR_dKzeta.item(), d2H_dKzeta2.item(), d2H_dKzeta_dKZ.item()],
-                [d2H_dKR_dKZ.item(), d2H_dKzeta_dKZ.item(), d2H_dKZ2.item()],
-            ]
-        )
-    )
+    dH_dR = dH["dH_dR"]
+    dH_dZ = dH["dH_dZ"]
+    dH_dKR = dH["dH_dKR"]
+    dH_dKzeta = dH["dH_dKzeta"]
+    dH_dKZ = dH["dH_dKZ"]
 
     d_Psi_d_tau = (
         -grad_grad_H
@@ -918,46 +160,4 @@ def beam_evolution_fun(
         - np.matmul(np.matmul(Psi_3D, gradK_gradK_H), Psi_3D)
     )
 
-    d_beam_parameters_d_tau = np.zeros_like(beam_parameters)
-
-    d_beam_parameters_d_tau[0] = dH_dKR  # d (q_R) / d tau
-    d_beam_parameters_d_tau[1] = dH_dKzeta  # d (q_zeta) / d tau
-
-    d_beam_parameters_d_tau[2] = dH_dKZ  # d (q_Z) / d tau
-    d_beam_parameters_d_tau[3] = -dH_dR  # d (K_R) / d tau
-    d_beam_parameters_d_tau[4] = -dH_dZ  # d (K_Z) / d tau
-
-    d_beam_parameters_d_tau[5] = np.real(d_Psi_d_tau[0, 0])  # d (Psi_RR) / d tau
-    d_beam_parameters_d_tau[6] = np.real(d_Psi_d_tau[1, 1])  # d (Psi_zetazeta) / d tau
-    d_beam_parameters_d_tau[7] = np.real(d_Psi_d_tau[2, 2])  # d (Psi_ZZ) / d tau
-    d_beam_parameters_d_tau[8] = np.real(d_Psi_d_tau[0, 1])  # d (Psi_Rzeta) / d tau
-    d_beam_parameters_d_tau[9] = np.real(d_Psi_d_tau[0, 2])  # d (Psi_RZ) / d tau
-    d_beam_parameters_d_tau[10] = np.real(d_Psi_d_tau[1, 2])  # d (Psi_zetaZ) / d tau
-
-    d_beam_parameters_d_tau[11] = np.imag(d_Psi_d_tau[0, 0])  # d (Psi_RR) / d tau
-    d_beam_parameters_d_tau[12] = np.imag(d_Psi_d_tau[1, 1])  # d (Psi_zetazeta) / d tau
-    d_beam_parameters_d_tau[13] = np.imag(d_Psi_d_tau[2, 2])  # d (Psi_ZZ) / d tau
-    d_beam_parameters_d_tau[14] = np.imag(d_Psi_d_tau[0, 1])  # d (Psi_Rzeta) / d tau
-    d_beam_parameters_d_tau[15] = np.imag(d_Psi_d_tau[0, 2])  # d (Psi_RZ) / d tau
-    d_beam_parameters_d_tau[16] = np.imag(d_Psi_d_tau[1, 2])  # d (Psi_zetaZ) / d tau
-
-    return d_beam_parameters_d_tau
-
-
-# ----------------------------------
-
-
-# def find_second_derivatives(K_magnitude,
-#                             K_magnitude_p_dKZ_1, K_magnitude_m_dKZ_1,
-#                             sin_theta_m_sq,
-#                             electron_density,
-#                             B_Total,
-#                             launch_angular_frequency, mode_flag,
-#                             delta_R, delta_Z, delta_K_R, delta_K_zeta, delta_K_Z):
-
-#     H_0 = find_H_numba(K_magnitude, electron_density, B_Total, sin_theta_m_sq, launch_angular_frequency, mode_flag)
-
-#     H_p_dZ = find_H_numba(K_magnitude, electron_density, B_Total, sin_theta_m_sq, launch_angular_frequency, mode_flag)
-
-
-#     return grad_grad_H, grad_gradK_H, gradK_gradK_H
+    return pack_beam_parameters(dH_dKR, dH_dKzeta, dH_dKZ, -dH_dR, -dH_dZ, d_Psi_d_tau)
diff --git a/scotty/fun_general.py b/scotty/fun_general.py
index c5fd2364..06fda8aa 100644
--- a/scotty/fun_general.py
+++ b/scotty/fun_general.py
@@ -198,10 +198,19 @@ def find_area_points(xs, ys, fraction_wanted):
 
 
 # Functions (Coordinate transformations)
+def freq_GHz_to_angular_frequency(freq_GHz: float) -> float:
+    """Convert frequency in GHz to angular frequency"""
+    return 2 * np.pi * 1e9 * freq_GHz
+
+
+def angular_frequency_to_wavenumber(angular_frequency: float) -> float:
+    """Convert angular frequency to wavenumber"""
+    return angular_frequency / constants.c
+
+
 def freq_GHz_to_wavenumber(freq_GHz: float) -> float:
     """Converts frequency in GHz to wavenumber"""
-    angular_frequency = 2 * np.pi * 1e9 * freq_GHz
-    return angular_frequency / constants.c
+    return angular_frequency_to_wavenumber(freq_GHz_to_angular_frequency(freq_GHz))
 
 
 def find_vec_lab_Cartesian(vec_lab, q_zeta):
@@ -485,74 +494,6 @@ def find_Booker_gamma(electron_density, B_Total, launch_angular_frequency):
 # ----------------------------------
 
 
-# Functions (beam tracing 2)
-def find_H(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    # For this functions to work,  the interpolation functions for
-    # electron density,  B_Total,  and psi  (poloidal flux) must be
-    # declared at some point before they are called
-
-    K_magnitude = np.sqrt(K_R**2 + (K_zeta / q_R) ** 2 + K_Z**2)
-    wavenumber_K0 = launch_angular_frequency / constants.c
-
-    poloidal_flux = interp_poloidal_flux(q_R, q_Z)
-    electron_density = find_density_1D(poloidal_flux)
-    B_R = np.squeeze(find_B_R(q_R, q_Z))
-    B_T = np.squeeze(find_B_T(q_R, q_Z))
-    B_Z = np.squeeze(find_B_Z(q_R, q_Z))
-
-    B_Total = np.sqrt(B_R**2 + B_T**2 + B_Z**2)
-    b_hat = np.array([B_R, B_T, B_Z]) / B_Total
-    K_hat = np.array([K_R, K_zeta / q_R, K_Z]) / K_magnitude
-    if np.size(q_R) == 1:
-        sin_theta_m_sq = (np.dot(b_hat, K_hat)) ** 2  # square of the mismatch angle
-    else:  # Vectorised version of find_H
-        b_hat = b_hat.T
-        K_hat = K_hat.T
-        sin_theta_m_sq = (
-            contract_special(b_hat, K_hat)
-        ) ** 2  # square of the mismatch angle
-
-    Booker_alpha = find_Booker_alpha(
-        electron_density, B_Total, sin_theta_m_sq, launch_angular_frequency
-    )
-    Booker_beta = find_Booker_beta(
-        electron_density, B_Total, sin_theta_m_sq, launch_angular_frequency
-    )
-    Booker_gamma = find_Booker_gamma(
-        electron_density, B_Total, launch_angular_frequency
-    )
-
-    # Due to numerical errors, sometimes H_discriminant ends up being a very small negative number
-    # That's why we take max(0, H_discriminant) in the sqrt
-
-    H = (K_magnitude / wavenumber_K0) ** 2 + (
-        Booker_beta
-        - mode_flag
-        * np.sqrt(
-            np.maximum(
-                np.zeros_like(Booker_beta),
-                (Booker_beta**2 - 4 * Booker_alpha * Booker_gamma),
-            )
-        )
-        # np.sqrt(Booker_beta**2 - 4*Booker_alpha*Booker_gamma)
-    ) / (2 * Booker_alpha)
-
-    return H
-
-
 def find_H_numba(
     K_magnitude,
     electron_density,
@@ -1467,5 +1408,23 @@ def genray_angles_from_mirror_angles(
 # ----------------------------------
 
 
-# Functions ()
-# ----------------------------------
+def make_array_3x3(array):
+    r"""Convert a 2x2 array into a 3x3 by appending zeros on the outside:
+
+    .. math::
+
+        \begin{pmatrix}
+            a & b \\
+            c & d \\
+        \end{pmatrix}
+        \Rightarrow
+        \begin{pmatrix}
+            a & b & 0 \\
+            c & d & 0 \\
+            0 & 0 & 0 \\
+        \end{pmatrix}
+    """
+    if array.shape != (2, 2):
+        raise ValueError(f"Expected array shape to be (2, 2), got {array.shape}")
+
+    return np.append(np.append(array, [[0, 0]], axis=0), [[0], [0], [0]], axis=1)
diff --git a/scotty/fun_mix.py b/scotty/fun_mix.py
deleted file mode 100644
index 38d20c6b..00000000
--- a/scotty/fun_mix.py
+++ /dev/null
@@ -1,323 +0,0 @@
-# -*- coding: utf-8 -*-
-"""
-Created on Fri Jun  8 10:44:34 2018
-
-Functions for finding derivatives of H using forward finite difference.
-
-@author: chenv
-Valerian Hongjie Hall-Chen
-valerian_hall-chen@ihpc.a-star.edu.sg
-
-Run in Python 3,  does not work in Python 2
-"""
-
-# import numpy as np
-# from scipy import constants as constants
-from scotty.fun_general import find_H
-from scotty.fun_FFD import find_dH_dR, find_dH_dZ  # \nabla H
-
-
-def find_d2H_dKR_dR(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_K_R,
-    delta_R,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    dH_dR_plus = find_dH_dR(
-        q_R,
-        q_Z,
-        K_R + delta_K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_R,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    dH_dR_minus = find_dH_dR(
-        q_R,
-        q_Z,
-        K_R - delta_K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_R,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKR_dR = (dH_dR_plus - dH_dR_minus) / (2 * delta_K_R)
-
-    return d2H_dKR_dR
-
-
-def find_d2H_dKR_dZ(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_K_R,
-    delta_Z,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    dH_dZ_plus = find_dH_dZ(
-        q_R,
-        q_Z,
-        K_R + delta_K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    dH_dZ_minus = find_dH_dZ(
-        q_R,
-        q_Z,
-        K_R - delta_K_R,
-        K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKR_dZ = (dH_dZ_plus - dH_dZ_minus) / (2 * delta_K_R)
-
-    return d2H_dKR_dZ
-
-
-def find_d2H_dKzeta_dR(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_K_zeta,
-    delta_R,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    dH_dR_plus = find_dH_dR(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta + delta_K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_R,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    dH_dR_minus = find_dH_dR(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta - delta_K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_R,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKzeta_dR = (dH_dR_plus - dH_dR_minus) / (2 * delta_K_zeta)
-
-    return d2H_dKzeta_dR
-
-
-def find_d2H_dKzeta_dZ(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_K_zeta,
-    delta_Z,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    dH_dZ_plus = find_dH_dZ(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta + delta_K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    dH_dZ_minus = find_dH_dZ(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta - delta_K_zeta,
-        K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKzeta_dZ = (dH_dZ_plus - dH_dZ_minus) / (2 * delta_K_zeta)
-
-    return d2H_dKzeta_dZ
-
-
-def find_d2H_dKZ_dR(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_K_Z,
-    delta_R,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    dH_dR_plus = find_dH_dR(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z + delta_K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_R,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    dH_dR_minus = find_dH_dR(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z - delta_K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_R,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKZ_dR = (dH_dR_plus - dH_dR_minus) / (2 * delta_K_Z)
-
-    return d2H_dKZ_dR
-
-
-def find_d2H_dKZ_dZ(
-    q_R,
-    q_Z,
-    K_R,
-    K_zeta,
-    K_Z,
-    launch_angular_frequency,
-    mode_flag,
-    delta_K_Z,
-    delta_Z,
-    interp_poloidal_flux,
-    find_density_1D,
-    find_B_R,
-    find_B_T,
-    find_B_Z,
-):
-    dH_dZ_plus = find_dH_dZ(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z + delta_K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    dH_dZ_minus = find_dH_dZ(
-        q_R,
-        q_Z,
-        K_R,
-        K_zeta,
-        K_Z - delta_K_Z,
-        launch_angular_frequency,
-        mode_flag,
-        delta_Z,
-        interp_poloidal_flux,
-        find_density_1D,
-        find_B_R,
-        find_B_T,
-        find_B_Z,
-    )
-    d2H_dKZ_dZ = (dH_dZ_plus - dH_dZ_minus) / (2 * delta_K_Z)
-
-    return d2H_dKZ_dZ
diff --git a/scotty/geometry.py b/scotty/geometry.py
index b9a98dbc..6cb97b37 100644
--- a/scotty/geometry.py
+++ b/scotty/geometry.py
@@ -27,6 +27,7 @@ class MagneticField(ABC):
 
     R_coord: FloatArray
     Z_coord: FloatArray
+    poloidalFlux_grid: FloatArray
 
     def B_R(self, q_R: ArrayLike, q_Z: ArrayLike) -> FloatArray:
         raise NotImplementedError
@@ -207,7 +208,7 @@ def _make_rect_spline(
         ky=interp_order,
         s=interp_smoothing,
     )
-    return lambda R, Z: spline(R, Z, grid=False)
+    return lambda q_R, q_Z: spline(q_R, q_Z, grid=False)
 
 
 class InterpolatedField(MagneticField):
diff --git a/scotty/hamiltonian.py b/scotty/hamiltonian.py
new file mode 100644
index 00000000..42a1dbd2
--- /dev/null
+++ b/scotty/hamiltonian.py
@@ -0,0 +1,449 @@
+from scotty.geometry import MagneticField
+from scotty.density_fit import DensityFitLike
+from scotty.fun_general import (
+    angular_frequency_to_wavenumber,
+    find_normalised_gyro_freq,
+    find_normalised_plasma_freq,
+    contract_special,
+)
+from scotty.typing import ArrayLike, FloatArray
+
+
+from dataclasses import dataclass, asdict
+import numpy as np
+from typing import Dict, Tuple
+
+
+class DielectricTensor:
+    r"""Calculates the components of the cold plasma dielectric tensor for a
+    wave with angular frequency :math:`\Omega`:
+
+    .. math::
+
+        \epsilon =
+          \begin{pmatrix}
+            \epsilon_{11}  & -i\epsilon_{12} & 0 \\
+            i\epsilon_{12} & \epsilon_{11}   & 0 \\
+            0            & 0             & \epsilon_{bb} \\
+          \end{pmatrix}
+
+    where:
+
+    .. math::
+
+        \begin{align}
+          \epsilon_{11} &= 1 - \frac{\Omega_{pe}^2}{\Omega^2 - \Omega_{ce}^2} \\
+          \epsilon_{12} &=
+            1 - \frac{\Omega_{pe}^2\Omega_{ce}}{\Omega(\Omega^2 - \Omega_{ce}^2)} \\
+          \epsilon_{bb} &= 1 - \frac{\Omega_{pe}^2}{\Omega^2} \\
+        \end{align}
+
+    The components of the dielectric tensor are calculated in the
+    :math:`(\hat{\mathbf{u}}_1, \hat{\mathbf{u}}_2, \hat{\mathbf{b}})` basis.
+    Hence, :math:`\epsilon_{11}`, :math:`\epsilon_{12}`, and
+    :math:`\epsilon_{bb}` correspond to the ``S``, ``D``, and ``P`` variables in
+    Stix, respectively. The notation used in this code is chosen to be
+    consistent with Hall-Chen, Parra, Hillesheim, PPCF 2022.
+
+    Parameters
+    ----------
+    electron_density:
+        Electron number density
+    angular_frequency:
+        Angular frequency of the beam
+    B_total:
+        Magnitude of the magnetic field
+
+    """
+
+    def __init__(
+        self,
+        electron_density: ArrayLike,
+        angular_frequency: float,
+        B_total: ArrayLike,
+    ):
+        _plasma_freq_2 = (
+            find_normalised_plasma_freq(electron_density, angular_frequency) ** 2
+        )
+        _gyro_freq = find_normalised_gyro_freq(B_total, angular_frequency)
+        _gyro_freq_2 = _gyro_freq**2
+
+        self._epsilon_bb = 1 - _plasma_freq_2
+        self._epsilon_11 = 1 - _plasma_freq_2 / (1 - _gyro_freq_2)
+        self._epsilon_12 = _plasma_freq_2 * _gyro_freq / (1 - _gyro_freq_2)
+
+    @property
+    def e_bb(self):
+        r"""The :math:`\epsilon_{bb}` component"""
+        return self._epsilon_bb
+
+    @property
+    def e_11(self):
+        r"""The :math:`\epsilon_{11}` component"""
+        return self._epsilon_11
+
+    @property
+    def e_12(self):
+        r"""The :math:`\epsilon_{12}` component"""
+        return self._epsilon_12
+
+
+Stencil = Dict[Tuple, float]
+
+FFD1_stencil: Stencil = {(0,): -3 / 2, (1,): 2, (2,): -0.5}
+"""Stencil for the first forward-difference"""
+FFD2_stencil: Stencil = {(0,): 2, (1,): -5, (2,): 4, (3,): -1}
+"""Stencil for the second forward-difference"""
+CFD1_stencil: Stencil = {(-1,): -0.5, (1,): 0.5}
+"""Stencil for the first central-difference"""
+CFD2_stencil: Stencil = {(-1,): 1, (0,): -2, (1,): 1}
+"""Stencil for the second central-difference"""
+FFD_FFD_stencil: Stencil = {
+    (0, 0): 2.25,
+    (0, 1): -3,
+    (0, 2): 0.75,
+    (1, 0): -3,
+    (1, 1): 4,
+    (1, 2): -1,
+    (2, 0): 0.75,
+    (2, 1): -1,
+    (2, 2): 0.25,
+}
+"""Stencil for the mixed second forward-difference"""
+FFD_CFD_stencil: Stencil = {
+    (0, -1): 0.25,
+    (0, 0): 1,
+    (0, 1): -1.25,
+    (1, 0): -2,
+    (1, 1): 2,
+    (2, -1): -0.25,
+    (2, 0): 1,
+    (2, 1): -0.75,
+}
+"""Stencil for the mixed second forward/central-difference"""
+CFD_CFD_stencil: Stencil = {
+    (1, 1): 0.25,
+    (1, -1): -0.25,
+    (-1, 1): -0.25,
+    (-1, -1): 0.25,
+}
+"""Stencil for the mixed second central-difference"""
+
+
+class Hamiltonian:
+    r"""Functor to evaluate derivatives of the Hamiltonian, :math:`H`, at a
+    given set of points.
+
+    ``Scotty`` calculates derivatives using a grid-free finite difference
+    approach. The Hamiltonian is evaluated at, essentially, an arbitrary set of
+    points around the location we wish to get the derivatives at. In practice we
+    define stencils as relative offsets from a central point, and the evaluation
+    points are the product of the spacing in a given direction with the stencil
+    offsets. By carefully choosing our stencils and evaluating all of the
+    derivatives at once, we can reuse evaluations of :math:`H` between
+    derivatives, saving a lot of computation.
+
+    The stencils are defined as a `dict` with a `tuple` of offsets as keys and
+    `float` weights as values. For example, the `CFD1_stencil`::
+
+        {(1,): 0.5, (-1,): -0.5}
+
+    defines the second-order first central-difference:
+
+    .. math::
+
+        f' = \frac{f(x + \delta_x) - f(x - \delta_x)}{2\delta_x}
+
+    The keys are tuples so that we can iterate over the offsets for the mixed
+    second derivatives.
+
+    The stencils have been chosen to maximise the reuse of Hamiltonian
+    evaluations without sacrificing accuracy.
+
+    Parameters
+    ----------
+    field
+        An object describing the magnetic field of the plasma
+    launch_angular_frequency
+        Angular frequency of the beam
+    mode_flag
+        Either ``+/-1``, used to determine which mode branch to use
+    density_fit
+        Function or ``Callable`` parameterising the density
+    delta_R
+        Finite difference spacing in the ``R`` direction
+    delta_Z
+        Finite difference spacing in the ``Z`` direction
+    delta_K_R
+        Finite difference spacing in the ``K_R`` direction
+    delta_K_zeta
+        Finite difference spacing in the ``K_zeta`` direction
+    delta_K_Z
+        Finite difference spacing in the ``K_Z`` direction
+
+    """
+
+    def __init__(
+        self,
+        field: MagneticField,
+        launch_angular_frequency: float,
+        mode_flag: int,
+        density_fit: DensityFitLike,
+        delta_R: float,
+        delta_Z: float,
+        delta_K_R: float,
+        delta_K_zeta: float,
+        delta_K_Z: float,
+    ):
+        self.field = field
+        self.angular_frequency = launch_angular_frequency
+        self.wavenumber_K0 = angular_frequency_to_wavenumber(launch_angular_frequency)
+        self.mode_flag = mode_flag
+        self.density = density_fit
+        self.spacings = {
+            "q_R": delta_R,
+            "q_Z": delta_Z,
+            "K_R": delta_K_R,
+            "K_zeta": delta_K_zeta,
+            "K_Z": delta_K_Z,
+        }
+
+    def __call__(
+        self,
+        q_R: ArrayLike,
+        q_Z: ArrayLike,
+        K_R: ArrayLike,
+        K_zeta: ArrayLike,
+        K_Z: ArrayLike,
+    ):
+        """Evaluate the Hamiltonian at the given coordinates
+
+        Parameters
+        ----------
+        q_R : ArrayLike
+        q_Z : ArrayLike
+        K_R : ArrayLike
+        K_zeta : ArrayLike
+        K_Z : ArrayLike
+
+        Returns
+        -------
+        ArrayLike
+
+        """
+
+        K_magnitude = np.sqrt(K_R**2 + (K_zeta / q_R) ** 2 + K_Z**2)
+        poloidal_flux = self.field.poloidal_flux(q_R, q_Z)
+        electron_density = self.density(poloidal_flux)
+        B_R = np.squeeze(self.field.B_R(q_R, q_Z))
+        B_T = np.squeeze(self.field.B_T(q_R, q_Z))
+        B_Z = np.squeeze(self.field.B_Z(q_R, q_Z))
+
+        B_total = np.sqrt(B_R**2 + B_T**2 + B_Z**2)
+        b_hat = np.array([B_R, B_T, B_Z]) / B_total
+        K_hat = np.array([K_R, K_zeta / q_R, K_Z]) / K_magnitude
+
+        # square of the mismatch angle
+        if np.size(q_R) == 1:
+            sin_theta_m_sq = np.dot(b_hat, K_hat) ** 2
+        else:  # Vectorised version of find_H
+            b_hat = b_hat.T
+            K_hat = K_hat.T
+            sin_theta_m_sq = contract_special(b_hat, K_hat) ** 2
+
+        epsilon = DielectricTensor(electron_density, self.angular_frequency, B_total)
+
+        Booker_alpha = (epsilon.e_bb * sin_theta_m_sq) + epsilon.e_11 * (
+            1 - sin_theta_m_sq
+        )
+        Booker_beta = (-epsilon.e_11 * epsilon.e_bb * (1 + sin_theta_m_sq)) - (
+            epsilon.e_11**2 - epsilon.e_12**2
+        ) * (1 - sin_theta_m_sq)
+        Booker_gamma = epsilon.e_bb * (epsilon.e_11**2 - epsilon.e_12**2)
+
+        H_discriminant = np.maximum(
+            np.zeros_like(Booker_beta),
+            (Booker_beta**2 - 4 * Booker_alpha * Booker_gamma),
+        )
+
+        return (K_magnitude / self.wavenumber_K0) ** 2 + (
+            Booker_beta - self.mode_flag * np.sqrt(H_discriminant)
+        ) / (2 * Booker_alpha)
+
+    def derivatives(
+        self,
+        q_R: ArrayLike,
+        q_Z: ArrayLike,
+        K_R: ArrayLike,
+        K_zeta: ArrayLike,
+        K_Z: ArrayLike,
+        second_order: bool = False,
+    ) -> Dict[str, ArrayLike]:
+        """Evaluate the first-order derivative in all directions at the given
+        point(s), and optionally the second-order ones too
+
+        Parameters
+        ----------
+        q_R : ArrayLike
+        q_Z : ArrayLike
+        K_R : ArrayLike
+        K_zeta : ArrayLike
+        K_Z : ArrayLike
+            Coordinates to evaluate the derivatives at
+        second_order : bool
+            If ``True``, also evaluate the second derivatives
+
+        """
+
+        @dataclass(frozen=True)
+        class CoordOffset:
+            """Helper class to store relative offsets in a dict"""
+
+            q_R: int = 0
+            q_Z: int = 0
+            K_R: int = 0
+            K_zeta: int = 0
+            K_Z: int = 0
+
+        # We cache Hamiltonian evaluations only during this function call
+        cache: Dict[CoordOffset, ArrayLike] = {}
+
+        # Capture the location we want the derivatives at
+        starts = {"q_R": q_R, "q_Z": q_Z, "K_R": K_R, "K_zeta": K_zeta, "K_Z": K_Z}
+
+        def apply_stencil(dims: Tuple[str, ...], stencil: Stencil) -> ArrayLike:
+            """Apply a stencil in set of directions
+
+            Parameters
+            ----------
+            dims
+                Tuple of direction names
+            stencil
+                Finite difference stencil
+            """
+
+            # Collect the relative spacings for the derivative dimensions
+            dim_spacings = [self.spacings[dim] for dim in dims]
+            # For second order derivatives, remove repeated dimensions
+            if len(dims) == 2 and dims[1] == dims[0]:
+                dims = (dims[0],)
+
+            result = np.zeros_like(starts["q_R"])
+
+            for stencil_offsets, weight in stencil.items():
+                coord_offsets = CoordOffset(**dict(zip(dims, stencil_offsets)))
+                try:
+                    # See if we've already evaluated H here...
+                    H_at_offset = cache[coord_offsets]
+                except KeyError:
+                    # ...if not, let's do so now
+                    offsets = asdict(coord_offsets)
+                    # Convert integer offsets to absolute positions
+                    coords = {
+                        dim: start + offsets[dim] * self.spacings[dim]
+                        for dim, start in starts.items()
+                    }
+                    # Evaluate the derivative at this offset and cache it
+                    H_at_offset = self(**coords)
+                    cache[coord_offsets] = H_at_offset
+
+                result += weight * H_at_offset
+            return result / np.prod(dim_spacings)
+
+        # We always compute the first order derivatives
+        derivatives = {
+            "dH_dR": apply_stencil(("q_R",), FFD1_stencil),
+            "dH_dZ": apply_stencil(("q_Z",), FFD1_stencil),
+            "dH_dKR": apply_stencil(("K_R",), CFD1_stencil),
+            "dH_dKzeta": apply_stencil(("K_zeta",), CFD1_stencil),
+            "dH_dKZ": apply_stencil(("K_Z",), CFD1_stencil),
+        }
+
+        if second_order:
+            second_derivatives = {
+                "d2H_dR2": apply_stencil(("q_R", "q_R"), FFD2_stencil),
+                "d2H_dZ2": apply_stencil(("q_Z", "q_Z"), FFD2_stencil),
+                "d2H_dKR2": apply_stencil(("K_R", "K_R"), CFD2_stencil),
+                "d2H_dKzeta2": apply_stencil(("K_zeta", "K_zeta"), CFD2_stencil),
+                "d2H_dKZ2": apply_stencil(("K_Z", "K_Z"), CFD2_stencil),
+                "d2H_dR_dZ": apply_stencil(("q_R", "q_Z"), FFD_FFD_stencil),
+                "d2H_dR_dKR": apply_stencil(("q_R", "K_R"), FFD_CFD_stencil),
+                "d2H_dR_dKzeta": apply_stencil(("q_R", "K_zeta"), FFD_CFD_stencil),
+                "d2H_dR_dKZ": apply_stencil(("q_R", "K_Z"), FFD_CFD_stencil),
+                "d2H_dZ_dKR": apply_stencil(("q_Z", "K_R"), FFD_CFD_stencil),
+                "d2H_dZ_dKzeta": apply_stencil(("q_Z", "K_zeta"), FFD_CFD_stencil),
+                "d2H_dZ_dKZ": apply_stencil(("q_Z", "K_Z"), FFD_CFD_stencil),
+                "d2H_dKR_dKZ": apply_stencil(("K_R", "K_Z"), CFD_CFD_stencil),
+                "d2H_dKR_dKzeta": apply_stencil(("K_R", "K_zeta"), CFD_CFD_stencil),
+                "d2H_dKzeta_dKZ": apply_stencil(("K_zeta", "K_Z"), CFD_CFD_stencil),
+            }
+            derivatives.update(second_derivatives)
+        return derivatives
+
+
+def hessians(dH: dict):
+    r"""Compute the elements of the Hessian of the Hamiltonian:
+
+    .. math::
+          \begin{gather}
+            \nabla \nabla H \\
+            \nabla_K \nabla H \\
+            \nabla_K \nabla_K H \\
+          \end{gather}
+
+    given a ``dict`` containing the second derivatives of the Hamiltonian (as
+    returned from `Hamiltonian.derivatives` with ``second_order=True``)
+    """
+
+    d2H_dR2 = dH["d2H_dR2"]
+    d2H_dZ2 = dH["d2H_dZ2"]
+    d2H_dKR2 = dH["d2H_dKR2"]
+    d2H_dKzeta2 = dH["d2H_dKzeta2"]
+    d2H_dKZ2 = dH["d2H_dKZ2"]
+    d2H_dR_dZ = dH["d2H_dR_dZ"]
+    d2H_dKR_dR = dH["d2H_dR_dKR"]
+    d2H_dKzeta_dR = dH["d2H_dR_dKzeta"]
+    d2H_dKZ_dR = dH["d2H_dR_dKZ"]
+    d2H_dKR_dZ = dH["d2H_dZ_dKR"]
+    d2H_dKzeta_dZ = dH["d2H_dZ_dKzeta"]
+    d2H_dKZ_dZ = dH["d2H_dZ_dKZ"]
+    d2H_dKR_dKZ = dH["d2H_dKR_dKZ"]
+    d2H_dKR_dKzeta = dH["d2H_dKR_dKzeta"]
+    d2H_dKzeta_dKZ = dH["d2H_dKzeta_dKZ"]
+
+    zeros = np.zeros_like(d2H_dR2)
+
+    def reshape(array: FloatArray):
+        """Such that shape is [points,3,3] instead of [3,3,points]"""
+        if array.ndim == 2:
+            return array
+        return np.moveaxis(np.squeeze(array), 2, 0)
+
+    grad_grad_H = reshape(
+        np.array(
+            [
+                [d2H_dR2, zeros, d2H_dR_dZ],
+                [zeros, zeros, zeros],
+                [d2H_dR_dZ, zeros, d2H_dZ2],
+            ]
+        )
+    )
+    gradK_grad_H = reshape(
+        np.array(
+            [
+                [d2H_dKR_dR, zeros, d2H_dKR_dZ],
+                [d2H_dKzeta_dR, zeros, d2H_dKzeta_dZ],
+                [d2H_dKZ_dR, zeros, d2H_dKZ_dZ],
+            ]
+        )
+    )
+    gradK_gradK_H = reshape(
+        np.array(
+            [
+                [d2H_dKR2, d2H_dKR_dKzeta, d2H_dKR_dKZ],
+                [d2H_dKR_dKzeta, d2H_dKzeta2, d2H_dKzeta_dKZ],
+                [d2H_dKR_dKZ, d2H_dKzeta_dKZ, d2H_dKZ2],
+            ]
+        )
+    )
+    return grad_grad_H, gradK_grad_H, gradK_gradK_H
diff --git a/scotty/launch.py b/scotty/launch.py
new file mode 100644
index 00000000..f1019d59
--- /dev/null
+++ b/scotty/launch.py
@@ -0,0 +1,387 @@
+from scotty.fun_general import (
+    make_array_3x3,
+    find_Psi_3D_lab,
+    find_Psi_3D_plasma,
+    find_inverse_2D,
+    find_K_lab_Cartesian,
+    find_q_lab_Cartesian,
+    find_K_lab,
+    find_d_poloidal_flux_dR,
+    find_d_poloidal_flux_dZ,
+    angular_frequency_to_wavenumber,
+)
+from scotty.hamiltonian import Hamiltonian
+from scotty.typing import FloatArray
+from scotty.geometry import MagneticField
+
+from typing import Optional
+
+import numpy as np
+from scipy.optimize import direct, root_scalar
+
+
+def launch_beam(
+    toroidal_launch_angle_Torbeam: float,
+    poloidal_launch_angle_Torbeam: float,
+    launch_beam_width: float,
+    launch_beam_curvature: float,
+    launch_position: FloatArray,
+    launch_angular_frequency: float,
+    field: MagneticField,
+    hamiltonian: Hamiltonian,
+    vacuumLaunch_flag: bool = True,
+    vacuum_propagation_flag: bool = True,
+    Psi_BC_flag: bool = True,
+    poloidal_flux_enter: Optional[float] = None,
+    delta_R: float = -1e-4,
+    delta_Z: float = 1e-4,
+):
+    r"""
+    Propagate the beam from its initial position at the antenna to
+    *just* inside the plasma.
+
+    Parameters
+    ----------
+    toroidal_launch_angle_Torbeam: float
+        Toroidal angle of antenna in TORBEAM convention
+    poloidal_launch_angle_Torbeam: float
+        Poloidal angle of antenna in TORBEAM convention
+    launch_beam_width: float
+        Width of the beam at launch
+    launch_beam_curvature: float
+        Curvatuve of the beam at launch
+    launch_position: FloatArray
+        Position of the antenna in cylindrical coordinates
+    launch_angular_frequency: float
+        Angular frequency of the beam at launch
+    field: MagneticField
+        Object describing the magnetic field of the plasma
+    vacuumLaunch_flag: bool
+        If ``True``, launch beam from vacuum, otherwise beam launch
+        position is inside the plasma already
+    vacuum_propagation_flag: bool
+        If ``True``, run solver from the launch position, and don't
+        use analytical vacuum propagation
+    Psi_BC_flag: bool
+        If ``True``, use matching boundary conditions at plasma entry
+        position, otherwise do no special treatment at plasma boundary
+    poloidal_flux_enter: Optional
+        Normalised poloidal flux label of plasma boundary. Required if
+        ``vacuum_propagation_flag`` is ``True``
+    delta_R: float
+        Finite difference spacing to use for ``R``
+    delta_Z: float
+        Finite difference spacing to use for ``Z``
+
+    Returns
+    -------
+    K_initial: FloatArray
+        Wavevector at plasma entry point
+    initial_position: FloatArray
+        Coordinates of entry point
+    launch_K: FloatArray
+        Wavevector at launch point
+    Psi_3D_lab_initial: FloatArray
+    Psi_3D_lab_launch: FloatArray
+    Psi_3D_lab_entry: FloatArray
+    Psi_3D_lab_entry_cartersian: FloatArray
+    distance_from_launch_to_entry: float
+
+    """
+
+    toroidal_launch_angle = np.deg2rad(toroidal_launch_angle_Torbeam)
+    poloidal_launch_angle = np.deg2rad(poloidal_launch_angle_Torbeam)
+
+    wavenumber_K0 = angular_frequency_to_wavenumber(launch_angular_frequency)
+
+    K_R_launch = (
+        -wavenumber_K0 * np.cos(toroidal_launch_angle) * np.cos(poloidal_launch_angle)
+    )
+    K_zeta_launch = (
+        -wavenumber_K0
+        * np.sin(toroidal_launch_angle)
+        * np.cos(poloidal_launch_angle)
+        * launch_position[0]
+    )
+    K_Z_launch = -wavenumber_K0 * np.sin(poloidal_launch_angle)
+    launch_K = np.array([K_R_launch, K_zeta_launch, K_Z_launch])
+
+    poloidal_rotation_angle = poloidal_launch_angle + (np.pi / 2)
+
+    Psi_w_beam_diagonal = (
+        wavenumber_K0 * launch_beam_curvature + 2j * launch_beam_width ** (-2)
+    )
+    Psi_w_beam_launch_cartersian = np.eye(2) * Psi_w_beam_diagonal
+
+    rotation_matrix_pol = np.array(
+        [
+            [np.cos(poloidal_rotation_angle), 0, np.sin(poloidal_rotation_angle)],
+            [0, 1, 0],
+            [-np.sin(poloidal_rotation_angle), 0, np.cos(poloidal_rotation_angle)],
+        ]
+    )
+
+    rotation_matrix_tor = np.array(
+        [
+            [np.cos(toroidal_launch_angle), np.sin(toroidal_launch_angle), 0],
+            [-np.sin(toroidal_launch_angle), np.cos(toroidal_launch_angle), 0],
+            [0, 0, 1],
+        ]
+    )
+
+    rotation_matrix = np.matmul(rotation_matrix_pol, rotation_matrix_tor)
+    rotation_matrix_inverse = np.transpose(rotation_matrix)
+
+    Psi_3D_beam_launch_cartersian = make_array_3x3(Psi_w_beam_launch_cartersian)
+    Psi_3D_lab_launch_cartersian = np.matmul(
+        rotation_matrix_inverse,
+        np.matmul(Psi_3D_beam_launch_cartersian, rotation_matrix),
+    )
+    Psi_3D_lab_launch = find_Psi_3D_lab(
+        Psi_3D_lab_launch_cartersian,
+        launch_position[0],
+        launch_position[1],
+        K_R_launch,
+        K_zeta_launch,
+    )
+
+    if not vacuum_propagation_flag:
+        return (
+            [K_R_launch, K_zeta_launch, K_Z_launch],
+            launch_position,
+            launch_K,
+            Psi_3D_lab_launch,
+            Psi_3D_lab_launch,
+            None,
+            np.full_like(Psi_3D_lab_launch, fill_value=np.nan),
+            None,
+        )
+
+    if poloidal_flux_enter is None:
+        raise ValueError(
+            "Missing required 'poloidal_flux_enter' for vacuum-launched beam"
+        )
+
+    Psi_w_beam_inverse_launch_cartersian = find_inverse_2D(Psi_w_beam_launch_cartersian)
+
+    entry_position = find_entry_point(
+        launch_position,
+        poloidal_launch_angle,
+        toroidal_launch_angle,
+        poloidal_flux_enter,
+        field,
+    )
+
+    distance_from_launch_to_entry = np.sqrt(
+        launch_position[0] ** 2
+        + entry_position[0] ** 2
+        - 2
+        * launch_position[0]
+        * entry_position[0]
+        * np.cos(entry_position[1] - launch_position[1])
+        + (launch_position[2] - entry_position[2]) ** 2
+    )
+
+    # Calculate entry parameters from launch parameters
+    # That is, find beam at start of plasma given its parameters at the antenna
+    K_lab_launch = np.array([K_R_launch, K_zeta_launch, K_Z_launch])
+    K_lab_Cartesian_launch = find_K_lab_Cartesian(K_lab_launch, launch_position)
+    K_lab_Cartesian_entry = K_lab_Cartesian_launch
+    entry_position_Cartesian = find_q_lab_Cartesian(entry_position)
+    K_lab_entry = find_K_lab(K_lab_Cartesian_entry, entry_position_Cartesian)
+
+    K_R_entry = K_lab_entry[0]  # K_R
+    K_zeta_entry = K_lab_entry[1]
+    K_Z_entry = K_lab_entry[2]  # K_z
+
+    Psi_w_beam_inverse_entry_cartersian = (
+        distance_from_launch_to_entry / (wavenumber_K0) * np.eye(2)
+        + Psi_w_beam_inverse_launch_cartersian
+    )
+    # 'entry' is still in vacuum, so the components of Psi along g are
+    # all 0 (since \nabla H = 0)
+    Psi_3D_beam_entry_cartersian = make_array_3x3(
+        find_inverse_2D(Psi_w_beam_inverse_entry_cartersian)
+    )
+    Psi_3D_lab_entry_cartersian = np.matmul(
+        rotation_matrix_inverse,
+        np.matmul(Psi_3D_beam_entry_cartersian, rotation_matrix),
+    )
+
+    # Convert to cylindrical coordinates
+    Psi_3D_lab_entry = find_Psi_3D_lab(
+        Psi_3D_lab_entry_cartersian,
+        entry_position[0],
+        entry_position[1],
+        K_R_entry,
+        K_zeta_entry,
+    )
+
+    # -------------------
+    # Find initial parameters in plasma
+    # -------------------
+    K_R_initial = K_R_entry
+    K_zeta_initial = K_zeta_entry
+    K_Z_initial = K_Z_entry
+    initial_position = entry_position
+    if Psi_BC_flag:  # Use BCs
+        dH = hamiltonian.derivatives(
+            initial_position[0],
+            initial_position[2],
+            K_R_initial,
+            K_zeta_initial,
+            K_Z_initial,
+        )
+
+        dH_dR_initial = dH["dH_dR"]
+        dH_dZ_initial = dH["dH_dZ"]
+        dH_dKR_initial = dH["dH_dKR"]
+        dH_dKzeta_initial = dH["dH_dKzeta"]
+        dH_dKZ_initial = dH["dH_dKZ"]
+        d_poloidal_flux_d_R_boundary = find_d_poloidal_flux_dR(
+            initial_position[0],
+            initial_position[2],
+            delta_R,
+            field.poloidal_flux,
+        )
+        d_poloidal_flux_d_Z_boundary = find_d_poloidal_flux_dZ(
+            initial_position[0],
+            initial_position[2],
+            delta_R,
+            field.poloidal_flux,
+        )
+
+        Psi_3D_lab_initial = find_Psi_3D_plasma(
+            Psi_3D_lab_entry,
+            dH_dKR_initial,
+            dH_dKzeta_initial,
+            dH_dKZ_initial,
+            dH_dR_initial,
+            dH_dZ_initial,
+            d_poloidal_flux_d_R_boundary,
+            d_poloidal_flux_d_Z_boundary,
+        )
+    else:  # Do not use BCs
+        Psi_3D_lab_initial = Psi_3D_lab_entry
+
+    return (
+        [K_R_initial, K_zeta_initial, K_Z_initial],
+        initial_position,
+        launch_K,
+        Psi_3D_lab_initial,
+        Psi_3D_lab_launch,
+        Psi_3D_lab_entry,
+        Psi_3D_lab_entry_cartersian,
+        distance_from_launch_to_entry,
+    )
+
+
+def find_entry_point(
+    launch_position: FloatArray,
+    poloidal_launch_angle: float,
+    toroidal_launch_angle: float,
+    poloidal_flux_enter: float,
+    field: MagneticField,
+    boundary_adjust: float = 1e-8,
+) -> FloatArray:
+    """Find the coordinates where the beam enters the plasma.
+
+    Arguments
+    ---------
+    launch_position:
+        Cartesian coordinates of the antenna (or cylindrical with
+        zeta=0)
+    poloidal_launch_angle:
+        Poloidal angle of the antenna (radians), clockwise from the
+        horizontal axis
+    toroidal_launch_angle:
+        Toroidal angle of the antenna (radians), anti-clockwise from
+        the negative X-axis
+    field:
+        Object describing the magnetic field geometry
+    boundary_adjust:
+        Step size used to ensure entry point is _just_ inside plasma
+
+    Returns
+    -------
+    Array with cylindrical coordinates of entry point
+    """
+
+    # We know that the plasma is contained entirely within ``field``'s
+    # (R, Z) grid, so the maximum distance the beam could possibly
+    # travel before hitting the plasma is when it's aimed at the
+    # top/bottom corner of the grid on the far side of the torus. This
+    # an overestimate, but it's only used to parameterize the beam
+    X_start, Y_start, Z_start = launch_position
+    X_length = abs(X_start) + field.R_coord.max()
+    Z_length = abs(Z_start) + field.Z_coord.max()
+    max_length = np.hypot(X_length, Z_length)
+
+    # TORBEAM antenna angles are anti-clockwise from negative X-axis,
+    # so we need to rotate the toroidal angle by pi. This will take
+    # care of the direction of the beam. The poloidal angle is also
+    # reversed from its usual sense, so we can just flip it
+    toroidal_launch_angle = toroidal_launch_angle + np.pi
+    poloidal_launch_angle = -poloidal_launch_angle
+
+    # We parameterise beam in a line normal to the antenna up to
+    # max_length, and we can be sure the beam will either hit the
+    # plasma or miss it entirely.
+    X_step = max_length * np.cos(toroidal_launch_angle) * np.cos(poloidal_launch_angle)
+    Y_step = max_length * np.sin(toroidal_launch_angle) * np.cos(poloidal_launch_angle)
+    Z_step = max_length * np.sin(poloidal_launch_angle)
+    step_array = np.array((X_step, Y_step, Z_step))
+
+    def beam_line(tau):
+        """Parameterised line in beam direction"""
+        return launch_position + tau * step_array
+
+    def cartesian_to_cylindrical(X, Y, Z):
+        """Cartesian to cylindrical coordinates"""
+        return np.sqrt(X**2 + Y**2), np.arctan2(Y, X), Z
+
+    def poloidal_flux_boundary_along_line(tau):
+        """Signed poloidal flux distance to plasma boundary"""
+        R, _, Z = cartesian_to_cylindrical(*beam_line(tau))
+        return field.poloidal_flux(R, Z) - poloidal_flux_enter
+
+    # Find the plasma boundary using a bracket search. We need the
+    # minimum so that the flux has different sign to our initial
+    # position for the bracket points. There could be one or two
+    # minima (corresponding to the two halves of the torus), we only
+    # care about the closest, so use the first result as a bound for
+    # the second. If there isn't a second, this will just find the
+    # first again
+    first_min = direct(
+        poloidal_flux_boundary_along_line, bounds=((0, 1),), len_tol=1e-3
+    )
+    minimum = direct(
+        poloidal_flux_boundary_along_line, bounds=((0, first_min.x),), len_tol=1e-3
+    )
+    # If the closest minimum is positive, then the beam never actually
+    # enters the plasma, and we should abort
+    if minimum.fun > 0:
+        R_miss, zeta_miss, Z_miss = cartesian_to_cylindrical(*beam_line(minimum.x))
+        miss_coords = f"(R={R_miss}, zeta={zeta_miss}, Z={Z_miss})"
+        raise RuntimeError(
+            f"Beam does not hit plasma. Closest point is at {miss_coords}, "
+            f"distance in poloidal flux to boundary={minimum.fun}"
+        )
+
+    # `minimum.x` is some point along the beam that's definitely
+    # inside the plasma. So any root between the start of the beam and
+    # this point is the closest entry point
+    boundary = root_scalar(poloidal_flux_boundary_along_line, bracket=(0, minimum.x))
+    if not boundary.converged:
+        raise RuntimeError(
+            f"Could not find plasma boundary, root finding failed with '{boundary.flag}'"
+        )
+
+    # The root might be just outside the plasma due to floating point
+    # errors, if so, take small steps until we're definitely inside
+    boundary_tau = boundary.root
+    R_boundary, _, Z_boundary = cartesian_to_cylindrical(*beam_line(boundary_tau))
+    if field.poloidal_flux(R_boundary, Z_boundary) > poloidal_flux_enter:
+        boundary_tau += boundary_adjust
+
+    return np.array(cartesian_to_cylindrical(*beam_line(boundary.root)))


Test patch:
diff --git a/tests/test_fun_general.py b/tests/test_fun_general.py
index 5edf7758..869e7ac8 100644
--- a/tests/test_fun_general.py
+++ b/tests/test_fun_general.py
@@ -2,6 +2,9 @@
     freq_GHz_to_wavenumber,
     find_nearest,
     read_floats_into_list_until,
+    find_Psi_3D_lab,
+    find_Psi_3D_lab_Cartesian,
+    make_array_3x3,
 )
 
 import io
@@ -61,3 +64,73 @@ def test_read_floats_into_list():
     npt.assert_allclose(x_coords, expected_x)
     npt.assert_allclose(z_coords, expected_z)
     npt.assert_allclose(B_r, expected_B_r)
+
+
+def random_symmetric_matrix():
+    array = np.random.random((3, 3))
+    array[0, 1] = array[1, 0]
+    array[0, 2] = array[2, 0]
+    array[1, 2] = array[2, 1]
+    return array
+
+
+def test_find_Psi_3D_lab_symmetry():
+    Psi_lab_Cartesian = random_symmetric_matrix()
+
+    Psi_lab_cylindrical = find_Psi_3D_lab(Psi_lab_Cartesian, 2, np.pi, 42, 64)
+
+    # result should be symmetric
+    assert Psi_lab_cylindrical[0, 1] == Psi_lab_cylindrical[1, 0]
+    assert Psi_lab_cylindrical[0, 2] == Psi_lab_cylindrical[2, 0]
+    assert Psi_lab_cylindrical[1, 2] == Psi_lab_cylindrical[2, 1]
+
+    # ZZ component unchanged
+    assert Psi_lab_cylindrical[2, 2].real == Psi_lab_Cartesian[2, 2]
+
+
+def test_convert_Psi_roundtrip():
+    Psi_original = random_symmetric_matrix()
+    q_R = 2
+    q_zeta = np.pi
+    K_R = 42
+    K_zeta = 64
+    Psi_cylindrical = find_Psi_3D_lab(Psi_original, q_R, q_zeta, K_R, K_zeta)
+    Psi_cartesian = find_Psi_3D_lab_Cartesian(Psi_cylindrical, q_R, q_zeta, K_R, K_zeta)
+
+    npt.assert_allclose(Psi_cartesian.real, Psi_original)
+    npt.assert_allclose(Psi_cartesian.imag, np.zeros_like(Psi_original))
+
+
+def test_Psi_cylindrical():
+    xx = 1
+    yy = 2
+    zz = 3
+    xy = 4
+    xz = 5
+    yz = 6
+    Psi_lab_Cartesian = np.array(
+        [
+            [xx, xy, xz],
+            [xy, yy, yz],
+            [xz, yz, zz],
+        ]
+    )
+
+    R = 2
+    Psi_lab_cylindrical = find_Psi_3D_lab(Psi_lab_Cartesian, R, np.pi, 0, 0)
+    expected = np.array(
+        [
+            [xx, xy * R, -xz],
+            [xy * R, yy * R**2, -yz * R],
+            [-xz, -yz * R, zz],
+        ]
+    )
+    npt.assert_allclose(Psi_lab_cylindrical, expected + 0j)
+
+
+def test_make_array_3x3():
+    A = np.array([[1, 2], [3, 4]])
+    B = make_array_3x3(A)
+
+    expected = np.array([[1, 2, 0], [3, 4, 0], [0, 0, 0]])
+    npt.assert_array_equal(B, expected)
diff --git a/tests/test_gradient.py b/tests/test_gradient.py
new file mode 100644
index 00000000..896026c7
--- /dev/null
+++ b/tests/test_gradient.py
@@ -0,0 +1,37 @@
+from scotty.fun_CFD import cfd_gradient
+
+
+import numpy as np
+
+
+def myfun(x, y):
+    return y**2 * (1 + np.sin(x))
+
+
+def test_gradient():
+    grad_x = cfd_gradient(myfun, x=0, y=3, directions="x", dx=1e-4)
+    assert np.isclose(grad_x, 9)
+
+    grad_y = cfd_gradient(myfun, x=0, y=3, directions="y", dx=1e-4)
+    assert np.isclose(grad_y, 6)
+
+
+def test_second_gradient():
+    grad_x = cfd_gradient(myfun, x=0, y=3, directions="x", dx=1e-4, d_order=2)
+    assert np.isclose(grad_x, 0, atol=1e-6)
+
+    grad_y = cfd_gradient(myfun, x=0, y=3, directions="y", dx=1e-4, d_order=2)
+    assert np.isclose(grad_y, 2)
+
+
+def test_gradient_xy():
+    grad_xy = cfd_gradient(
+        myfun, x=0, y=3, directions=("x", "y"), dx=(1e-4, 1e-4), d_order=2
+    )
+    assert np.isclose(grad_xy, 6)
+
+    # Shouldn't depend on order of directions
+    grad_yx = cfd_gradient(
+        myfun, x=0, y=3, directions=("y", "x"), dx=(1e-4, 1e-4), d_order=2
+    )
+    assert np.isclose(grad_yx, 6)
diff --git a/tests/test_hamiltonian.py b/tests/test_hamiltonian.py
new file mode 100644
index 00000000..ea915a9b
--- /dev/null
+++ b/tests/test_hamiltonian.py
@@ -0,0 +1,178 @@
+from scotty.hamiltonian import Hamiltonian, hessians
+from scotty.fun_general import freq_GHz_to_angular_frequency
+from scotty.init_bruv import get_parameters_for_Scotty
+from scotty.beam_me_up import create_magnetic_geometry
+
+import numpy as np
+from numpy.testing import assert_allclose
+
+import pytest
+
+
+# Wavenumbers used in the tests
+k_q_R = 2.1
+k_q_Z = 3.2
+k_K_R = 4.3
+k_K_zeta = 5.4
+k_K_Z = 6.5
+
+
+class FakeHamiltonian(Hamiltonian):
+    def __init__(
+        self,
+        delta_R: float,
+        delta_Z: float,
+        delta_K_R: float,
+        delta_K_zeta: float,
+        delta_K_Z: float,
+    ):
+        self.spacings = {
+            "q_R": delta_R,
+            "q_Z": delta_Z,
+            "K_R": delta_K_R,
+            "K_zeta": delta_K_zeta,
+            "K_Z": delta_K_Z,
+        }
+
+    def __call__(self, q_R, q_Z, K_R, K_zeta, K_Z):
+        return (
+            np.exp(k_q_R * q_R)
+            * np.exp(k_q_Z * q_Z)
+            * np.exp(k_K_R * K_R)
+            * np.exp(k_K_zeta * K_zeta)
+            * np.exp(k_K_Z * K_Z)
+        )
+
+
+def test_hamiltonian_derivatives():
+    H = FakeHamiltonian(1e-3, 1e-3, 1e-4, 1e-4, 1e-4)
+    dH = H.derivatives(1.2, 2.3, 3.4, 4.5, 5.6, second_order=True)
+
+    H0 = H(1.2, 2.3, 3.4, 4.5, 5.6)
+
+    dH_dR = dH["dH_dR"] / H0
+    dH_dZ = dH["dH_dZ"] / H0
+    dH_dKR = dH["dH_dKR"] / H0
+    dH_dKzeta = dH["dH_dKzeta"] / H0
+    dH_dKZ = dH["dH_dKZ"] / H0
+    d2H_dR2 = dH["d2H_dR2"] / H0
+    d2H_dZ2 = dH["d2H_dZ2"] / H0
+    d2H_dKR2 = dH["d2H_dKR2"] / H0
+    d2H_dKzeta2 = dH["d2H_dKzeta2"] / H0
+    d2H_dKZ2 = dH["d2H_dKZ2"] / H0
+    d2H_dR_dZ = dH["d2H_dR_dZ"] / H0
+    d2H_dKR_dR = dH["d2H_dR_dKR"] / H0
+    d2H_dKzeta_dR = dH["d2H_dR_dKzeta"] / H0
+    d2H_dKZ_dR = dH["d2H_dR_dKZ"] / H0
+    d2H_dKR_dZ = dH["d2H_dZ_dKR"] / H0
+    d2H_dKzeta_dZ = dH["d2H_dZ_dKzeta"] / H0
+    d2H_dKZ_dZ = dH["d2H_dZ_dKZ"] / H0
+    d2H_dKR_dKZ = dH["d2H_dKR_dKZ"] / H0
+    d2H_dKR_dKzeta = dH["d2H_dKR_dKzeta"] / H0
+    d2H_dKzeta_dKZ = dH["d2H_dKzeta_dKZ"] / H0
+
+    assert np.isclose(dH_dR, k_q_R)
+    assert np.isclose(dH_dZ, k_q_Z)
+    assert np.isclose(dH_dKR, k_K_R)
+    assert np.isclose(dH_dKzeta, k_K_zeta)
+    assert np.isclose(dH_dKZ, k_K_Z)
+    assert np.isclose(d2H_dR2, k_q_R**2)
+    assert np.isclose(d2H_dZ2, k_q_Z**2)
+    assert np.isclose(d2H_dKR2, k_K_R**2)
+    assert np.isclose(d2H_dKzeta2, k_K_zeta**2)
+    assert np.isclose(d2H_dKZ2, k_K_Z**2)
+    assert np.isclose(d2H_dR_dZ, k_q_R * k_q_Z)
+    assert np.isclose(d2H_dKR_dR, k_K_R * k_q_R)
+    assert np.isclose(d2H_dKzeta_dR, k_K_zeta * k_q_R)
+    assert np.isclose(d2H_dKZ_dR, k_K_Z * k_q_R)
+    assert np.isclose(d2H_dKR_dZ, k_K_R * k_q_Z)
+    assert np.isclose(d2H_dKzeta_dZ, k_K_zeta * k_q_Z)
+    assert np.isclose(d2H_dKZ_dZ, k_K_Z * k_q_Z)
+    assert np.isclose(d2H_dKR_dKZ, k_K_R * k_K_Z)
+    assert np.isclose(d2H_dKR_dKzeta, k_K_R * k_K_zeta)
+    assert np.isclose(d2H_dKzeta_dKZ, k_K_zeta * k_K_Z)
+
+
+def test_hamiltonian_hessians():
+    H = FakeHamiltonian(1e-3, 1e-3, 1e-4, 1e-4, 1e-4)
+    dH = H.derivatives(1.2, 2.3, 3.4, 4.5, 5.6, second_order=True)
+    H0 = H(1.2, 2.3, 3.4, 4.5, 5.6)
+
+    grad_grad_H, gradK_grad_H, gradK_gradK_H = hessians(dH)
+
+    grad_grad_H_expected = np.array(
+        [[k_q_R**2, 0, k_q_R * k_q_Z], [0, 0, 0], [k_q_R * k_q_Z, 0, k_q_Z**2]]
+    )
+    gradK_grad_H_expected = np.array(
+        [
+            [k_K_R * k_q_R, 0, k_K_R * k_q_Z],
+            [k_K_zeta * k_q_R, 0, k_K_zeta * k_q_Z],
+            [k_K_Z * k_q_R, 0, k_K_Z * k_q_Z],
+        ]
+    )
+    gradK_gradK_H_expected = np.array(
+        [
+            [k_K_R**2, k_K_R * k_K_zeta, k_K_R * k_K_Z],
+            [k_K_zeta * k_K_R, k_K_zeta**2, k_K_zeta * k_K_Z],
+            [k_K_Z * k_K_R, k_K_Z * k_K_zeta, k_K_Z**2],
+        ]
+    )
+    assert_allclose(grad_grad_H / H0, grad_grad_H_expected, rtol=1e-5, atol=1e-4)
+    assert_allclose(gradK_grad_H / H0, gradK_grad_H_expected, rtol=1e-5, atol=1e-4)
+    assert_allclose(gradK_gradK_H / H0, gradK_gradK_H_expected, rtol=1e-5, atol=1e-4)
+
+
+@pytest.mark.parametrize(
+    ("derivative", "expected"),
+    (
+        ("dH_dR", k_q_R),
+        ("d2H_dR2", k_q_R**2),
+        ("d2H_dR_dZ", k_q_R * k_q_Z),
+        ("d2H_dKZ2", k_K_Z**2),
+        ("d2H_dKzeta_dKZ", k_K_zeta * k_K_Z),
+        ("d2H_dR_dKzeta", k_K_zeta * k_q_R),
+    ),
+)
+def test_hamiltonian_derivatives_scaling(derivative, expected):
+    """Check that error scaling of derivatives is second order"""
+    dx = 1e-3
+    dx4 = dx / 4
+    H = FakeHamiltonian(dx, dx, dx, dx, dx)
+    dH_0 = H.derivatives(1.2, 2.3, 3.4, 4.5, 5.6, second_order=True)
+
+    dH_1 = FakeHamiltonian(dx4, dx4, dx4, dx4, dx4).derivatives(
+        1.2, 2.3, 3.4, 4.5, 5.6, second_order=True
+    )
+
+    H0 = H(1.2, 2.3, 3.4, 4.5, 5.6)
+
+    error_0 = np.abs(dH_0[derivative] / H0 - expected)
+    error_1 = np.abs(dH_1[derivative] / H0 - expected)
+
+    err_diff = np.log(error_0 / error_1) / np.log(4)
+    assert np.isclose(err_diff, 2, rtol=1e-2)
+
+
+def test_golden():
+    """Check that the actual Hamiltonian is correct
+
+    Very basic golden answer test, could do with analytic answer
+    """
+    kwargs_dict = get_parameters_for_Scotty("DBS_synthetic")
+    kwargs_dict["find_B_method"] = "unit-tests"
+    field = create_magnetic_geometry(**kwargs_dict)
+    density = kwargs_dict["density_fit_method"]
+    angular_frequency = freq_GHz_to_angular_frequency(kwargs_dict["launch_freq_GHz"])
+    H = Hamiltonian(
+        field,
+        angular_frequency,
+        kwargs_dict["mode_flag"],
+        density,
+        1e-3,
+        1e-3,
+        1e-3,
+        1e-3,
+        1e-3,
+    )
+    expected = -0.2632447148279265
+    assert np.isclose(H(1.75, 0.1, 1, 1, 1), expected)
diff --git a/tests/test_simple_golden.py b/tests/test_simple_golden.py
index bdca806e..6d0ca98c 100644
--- a/tests/test_simple_golden.py
+++ b/tests/test_simple_golden.py
@@ -1,14 +1,26 @@
-from scotty.beam_me_up import beam_me_up, create_magnetic_geometry
+from scotty.beam_me_up import (
+    beam_me_up,
+    create_magnetic_geometry,
+    make_density_fit,
+)
 from scotty.init_bruv import get_parameters_for_Scotty
 from scotty.torbeam import Torbeam
 from scotty.geometry import CircularCrossSectionField
+from scotty.fun_general import (
+    freq_GHz_to_angular_frequency,
+    angular_frequency_to_wavenumber,
+)
+from scotty.hamiltonian import Hamiltonian
+from scotty.launch import find_entry_point, launch_beam
 
 import json
 import numpy as np
-from numpy.testing import assert_allclose
+from numpy.testing import assert_allclose, assert_array_equal
 import pathlib
 import pytest
 
+# Print more of arrays in failed tests
+np.set_printoptions(linewidth=120, threshold=100)
 
 CUTOFF_INDEX = 5
 
@@ -187,7 +199,7 @@
     ),
     "normalised_plasma_freqs": np.array(
         [
-            0.01388158,
+            0.0,
             0.59773859,
             0.75915389,
             0.83777097,
@@ -236,8 +248,8 @@ def synthetic_args():
         minor_radius_a=kwargs_dict["minor_radius_a"],
         B_T_axis=kwargs_dict["B_T_axis"],
         B_p_a=kwargs_dict["B_p_a"],
-        R_points=100,
-        Z_points=100,
+        R_points=256,
+        Z_points=256,
         grid_buffer_factor=1.1,
     )
     x_meshgrid, z_meshgrid = np.meshgrid(field.R_coord, field.Z_coord, indexing="ij")
@@ -363,7 +375,7 @@ def UDA_saved(path: pathlib.Path):
     kwargs_dict["find_B_method"] = "UDA_saved"
     kwargs_dict["magnetic_data_path"] = path / "12345_equilibrium_data.npz"
     kwargs_dict["shot"] = None
-    kwargs_dict["delta_R"] = 0.0001
+    kwargs_dict["delta_R"] = -0.0001
     kwargs_dict["delta_Z"] = 0.0001
     return kwargs_dict
 
@@ -400,7 +412,7 @@ def UDA_saved_MAST_U(path: pathlib.Path):
     kwargs_dict["find_B_method"] = "UDA_saved"
     kwargs_dict["magnetic_data_path"] = path
     kwargs_dict["shot"] = shot
-    kwargs_dict["delta_R"] = 0.0001
+    kwargs_dict["delta_R"] = -0.0001
     kwargs_dict["delta_Z"] = 0.0001
     kwargs_dict["equil_time"] = time[1]
     return kwargs_dict
@@ -438,18 +450,23 @@ def test_integrated(tmp_path, generator):
         output = dict(f)
 
     for key, value in EXPECTED.items():
-        assert_allclose(output[key], value, rtol=1e-2, err_msg=key)
+        assert_allclose(output[key], value, rtol=1e-2, atol=1e-2, err_msg=key)
 
     K_magnitude = np.hypot(output["K_R_array"], output["K_Z_array"])
     assert K_magnitude.argmin() == CUTOFF_INDEX
 
-    assert_allclose(output["Psi_3D_output"][0, ...], PSI_START_EXPECTED, rtol=1.8e-2)
     assert_allclose(
-        output["Psi_3D_output"][CUTOFF_INDEX, ...], PSI_CUTOFF_EXPECTED, rtol=1e-2
+        output["Psi_3D_output"][0, ...], PSI_START_EXPECTED, rtol=1e-2, atol=0.1
+    )
+    assert_allclose(
+        output["Psi_3D_output"][CUTOFF_INDEX, ...],
+        PSI_CUTOFF_EXPECTED,
+        rtol=1e-2,
+        atol=0.1,
+    )
+    assert_allclose(
+        output["Psi_3D_output"][-1, ...], PSI_FINAL_EXPECTED, rtol=1.8e-2, atol=0.1
     )
-    # Slightly larger tolerance here, likely due to floating-point
-    # precision of file-based input
-    assert_allclose(output["Psi_3D_output"][-1, ...], PSI_FINAL_EXPECTED, rtol=2e-2)
 
 
 @pytest.mark.parametrize(
@@ -524,3 +541,220 @@ def test_EFIT_geometry(tmp_path):
     assert total_sign[0, -1] == 1, "Top right"
     assert total_sign[-1, 0] == 1, "Bottom left"
     assert total_sign[-1, -1] == -1, "Bottom right"
+
+
+@pytest.mark.parametrize(
+    "generator",
+    [
+        pytest.param(simple, id="simple"),
+        pytest.param(ne_dat_file, id="density-fit-file"),
+        pytest.param(torbeam_file, id="torbeam-file"),
+        pytest.param(omfit_json, id="omfit-file"),
+        pytest.param(npz_file, id="test-file"),
+        pytest.param(npz_notime_file, id="test_notime-file"),
+        # pytest.param(UDA_saved, id="UDA-saved-file"),
+        # pytest.param(UDA_saved_MAST_U, id="UDA-saved-MAST-U-file"),
+    ],
+)
+def test_launch_golden_answer(tmp_path, generator):
+    args = generator(tmp_path)
+    field = create_magnetic_geometry(**args)
+    ne_data_path = args.get("ne_data_path", None)
+    if ne_data_path:
+        ne_filename = ne_data_path / "ne.dat"
+        parameters = [ne_filename, 5, 0]
+    else:
+        parameters = None
+        ne_filename = None
+
+    density_fit = make_density_fit(
+        args["density_fit_method"],
+        args["poloidal_flux_enter"],
+        parameters,
+        ne_filename,
+    )
+
+    launch_angular_frequency = freq_GHz_to_angular_frequency(args["launch_freq_GHz"])
+
+    hamiltonian = Hamiltonian(
+        field,
+        launch_angular_frequency,
+        mode_flag=args["mode_flag"],
+        density_fit=density_fit,
+        delta_R=-1e-3,
+        delta_Z=1e-3,
+        delta_K_R=0.1,
+        delta_K_zeta=0.1,
+        delta_K_Z=0.1,
+    )
+    (
+        K_initial,
+        initial_position,
+        launch_K,
+        Psi_3D_lab_initial,
+        Psi_3D_lab_launch,
+        Psi_3D_lab_entry,
+        Psi_3D_lab_entry_cartersian,
+        distance_from_launch_to_entry,
+    ) = launch_beam(
+        field=field,
+        hamiltonian=hamiltonian,
+        toroidal_launch_angle_Torbeam=args["toroidal_launch_angle_Torbeam"],
+        poloidal_launch_angle_Torbeam=args["poloidal_launch_angle_Torbeam"],
+        launch_beam_width=args["launch_beam_width"],
+        launch_beam_curvature=args["launch_beam_curvature"],
+        launch_position=args["launch_position"],
+        launch_angular_frequency=launch_angular_frequency,
+        vacuum_propagation_flag=args["vacuum_propagation_flag"],
+        Psi_BC_flag=args["Psi_BC_flag"],
+        poloidal_flux_enter=args["poloidal_flux_enter"],
+    )
+
+    expected_K_initial = np.array([-1146.4000699962507, -0.0, -120.4915026654739])
+    expected_initial_position = np.array([1.99382564, -0.0, -0.07804514])
+    expected_launch_K = np.array([-1146.40007, -0.0, -120.49150267])
+    # Note that R-zeta, Z-zeta terms are not actually zero, but should be
+    expected_Psi_3D_lab_initial = np.array(
+        [
+            [-2.473578e03 + 1.195578e01j, 0.0 + 0.0j, 3.405239e02 - 1.137526e02j],
+            [0.0 + 0.0j, 4.245892e03 + 4.350010e03j, 0.0 + 0.0j],
+            [3.405239e02 - 1.137526e02j, 0.0 + 0.0j, 4.257764e02 + 1.082292e03j],
+        ]
+    )
+    expected_Psi_3D_lab_launch = np.array(
+        [
+            [-3.1486979 + 13.65774954j, 0.0 + 0.0j, 29.9578594 - 129.94480676j],
+            [0.0 + 0.0j, 1037.08121046 + 8365.71125j, 0.0 + 0.0j],
+            [29.9578594 - 129.94480676j, 0.0 + 0.0j, -285.02999262 + 1236.34225046j],
+        ]
+    )
+    expected_Psi_3D_lab_entry = np.array(
+        [
+            [5.38751589 + 11.95597601j, 0.0 + 0.0j, -51.25878964 - 113.75351313j],
+            [0.0 + 0.0j, 4245.89207422 + 4350.0100231j, 0.0 + 0.0j],
+            [-51.25878964 - 113.75351313j, 0.0 + 0.0j, 487.69480611 + 1082.29238188j],
+        ]
+    )
+    expected_Psi_3D_lab_entry_cartersian = np.array(
+        [
+            [5.38751589 + 11.95597601j, 0.0 + 0.0j, -51.25878964 - 113.75351313j],
+            [0.0 + 0.0j, 493.082322 + 1094.24835789j, 0.0 + 0.0j],
+            [-51.25878964 - 113.75351313j, 0.0 + 0.0j, 487.69480611 + 1082.29238188j],
+        ]
+    )
+    expected_distance_from_launch_to_entry = 0.5964417281350541
+    expected_Psi_3D = np.array(
+        [
+            [-3.1486979 + 13.65774954j, 0.0 + 0.0j, 29.9578594 - 129.94480676j],
+            [0.0 + 0.0j, 1037.08121046 + 8365.71125j, 0.0 + 0.0j],
+            [29.9578594 - 129.94480676j, 0.0 + 0.0j, -285.02999262 + 1236.34225046j],
+        ]
+    )
+
+    tol = 1e-3
+    assert_allclose(Psi_3D_lab_launch, expected_Psi_3D, tol, tol)
+    assert_allclose(K_initial, expected_K_initial, tol, tol)
+    assert_allclose(initial_position, expected_initial_position, tol, tol)
+    assert_allclose(launch_K, expected_launch_K, tol, tol)
+    assert_allclose(Psi_3D_lab_launch, expected_Psi_3D_lab_launch, tol, tol)
+    assert_allclose(Psi_3D_lab_entry, expected_Psi_3D_lab_entry, tol, tol)
+    assert_allclose(
+        Psi_3D_lab_entry_cartersian, expected_Psi_3D_lab_entry_cartersian, tol, tol
+    )
+    assert_allclose(
+        distance_from_launch_to_entry,
+        expected_distance_from_launch_to_entry,
+        tol,
+        tol,
+    )
+    assert_allclose(launch_K, expected_launch_K, tol, tol)
+    assert_allclose(Psi_3D_lab_launch, expected_Psi_3D_lab_launch, tol, tol)
+    assert_allclose(Psi_3D_lab_entry, expected_Psi_3D_lab_entry, tol, tol)
+    assert_allclose(
+        Psi_3D_lab_entry_cartersian, expected_Psi_3D_lab_entry_cartersian, tol, tol
+    )
+    # Larger atol due to some small values
+    assert_allclose(Psi_3D_lab_initial, expected_Psi_3D_lab_initial, tol, atol=0.1)
+
+
+def launch_parameters(start_point, end_point_poloidal_coords):
+    kwargs_dict = simple(None)
+    rho_end, theta_end, zeta_end = end_point_poloidal_coords
+    R_axis = kwargs_dict["R_axis"]
+    minor_radius = kwargs_dict["minor_radius_a"]
+
+    R_end = R_axis + rho_end * minor_radius * np.cos(theta_end)
+    Z_end = rho_end * minor_radius * np.sin(theta_end)
+    X_end = R_end * np.cos(zeta_end)
+    Y_end = R_end * np.sin(zeta_end)
+
+    X_start, Y_start, Z_start = start_point
+    phi_t = np.arctan2(Y_end - Y_start, X_end - X_start) - np.pi
+    poloidal_length = np.sqrt((Y_end - Y_start) ** 2 + (X_end - X_start) ** 2)
+    phi_p = np.arctan2(Z_end - Z_start, poloidal_length)
+    expected_entry_cartesian = (R_end, zeta_end, Z_end)
+    return start_point, -phi_p, phi_t, expected_entry_cartesian
+
+
+@pytest.mark.parametrize(
+    "generator",
+    [
+        pytest.param(simple, id="simple"),
+        pytest.param(torbeam_file, id="torbeam-file"),
+        pytest.param(npz_file, id="test-file"),
+        pytest.param(UDA_saved, id="UDA-saved-file"),
+    ],
+)
+@pytest.mark.parametrize(
+    (
+        "launch_position",
+        "poloidal_launch_angle",
+        "toroidal_launch_angle",
+        "expected_entry",
+    ),
+    (
+        pytest.param([2.5, 0, 0], 0, 0, [2.0, 0.0, 0], id="outboard-midplane"),
+        pytest.param([0.5, 0, 0], 0, np.pi, [1.0, 0, 0], id="inboard-midplane"),
+        pytest.param([1.5, 0, 1], np.pi / 2, 0, [1.5, 0, 0.5], id="top"),
+        pytest.param([1.5, 0, -1], -np.pi / 2, 0, [1.5, 0, -0.5], id="bottom"),
+        pytest.param(
+            [2, 0, 0.5],
+            np.pi / 4,
+            0,
+            [1.5 + np.sqrt(0.5) / 2, 0, np.sqrt(0.5) / 2],
+            id="top-right",
+        ),
+        pytest.param(
+            [2, 0, -0.5],
+            -np.pi / 4,
+            0,
+            [1.5 + np.sqrt(0.5) / 2, 0, -np.sqrt(0.5) / 2],
+            id="bottom-right",
+        ),
+        pytest.param(
+            *launch_parameters((2.5, 0.0, -0.1), (1, np.pi / 4, np.pi / 8)),
+            id="steep-angle",
+        ),
+    ),
+)
+def test_find_entry_point(
+    tmp_path,
+    generator,
+    launch_position,
+    poloidal_launch_angle,
+    toroidal_launch_angle,
+    expected_entry,
+):
+    args = generator(tmp_path)
+    field = create_magnetic_geometry(**args)
+    poloidal_flux_enter = args["poloidal_flux_enter"]
+
+    entry_position = find_entry_point(
+        launch_position,
+        poloidal_launch_angle,
+        toroidal_launch_angle,
+        poloidal_flux_enter,
+        field,
+    )
+
+    assert_allclose(entry_position, expected_entry, 1e-6, 1e-6)