astronomix.analysis_helpers.jacobians module

astronomix.analysis_helpers.jacobians module#

Helpers for computing Jacobians of the simulator.

For a base state that depends only on y, the streamwise (x) direction is translation-invariant, so different x-Fourier modes decouple in the linearised problem. These helpers assemble the dense Jacobian of a single +kx Fourier block — either of the spatial RHS or of the time-integrated simulator — column-by-column in small batches, avoiding the cost of forming the full grid-sized Jacobian.

astronomix.analysis_helpers.jacobians.single_xmode_jacobian2Dt(primitive_state_unperturbed, config, params, registered_variables, helper_data, wavelength, assembly_batch_size=4)[source]#

Return the dense +kx Fourier-block Jacobian of the time-integrated simulator.

This is the time-integration analogue of single_xmode_rhs_jacobian2D(): instead of linearising the instantaneous RHS R(U) it linearises the full time integration over [t_start, t_end]. Because the time map amplifies growing modes, this can bring out the eigenmodes with large Re(\lambda) more cleanly than the RHS Jacobian. The complex +kx block is assembled column-by-column from real JVPs exactly as in the RHS variant.

Parameters:

primitive_state_unperturbed – Primitive base state to linearise about.
config – Simulation config. The x-boundary should be periodic.
params – Simulation parameters.
registered_variables – Variable registry.
helper_data – Grid/helper data containing cell centers.
wavelength – Physical wavelength of the x-Fourier perturbation. Must be commensurate with the x-domain length.
assembly_batch_size – Number of Jacobian columns to compute per JVP batch. Lower this if GPU memory is tight.

Returns:

Complex array of shape (nvar * Ny, nvar * Ny), the raw +kx Fourier-block tangent operator of the time-integrated simulator.

Return type:

J_complex

astronomix.analysis_helpers.jacobians.single_xmode_rhs_jacobian2D(primitive_state_unperturbed, config, params, registered_variables, helper_data, wavelength, assembly_batch_size=4)[source]#

Return the dense +kx Fourier-block Jacobian of the simulator RHS.

The full semi-discrete RHS is

dU/dt = R(U),

where U is the conserved state on the full 2D grid. The full Jacobian dR/dU has size

(nvar * Nx * Ny) x (nvar * Nx * Ny),

which is usually too large to form explicitly.

If the unperturbed state U0 depends only on y, the x-direction is translation-invariant and different streamwise Fourier modes decouple in the linearized problem. We restrict the perturbation to one Fourier mode,

dU(x, y) = qhat(y) exp(i kx x),

with

qhat(y) in C^(nvar * Ny).

The returned matrix represents

L_k qhat = P_k J[ qhat(y) exp(i kx x) ],

where J = dR/dU evaluated at U0 and P_k projects back onto the +kx Fourier mode.

A complex Fourier amplitude is just a compact representation of sine and cosine perturbations:

Re[qhat exp(i kx x)] = Re(qhat) cos(kx x) - Im(qhat) sin(kx x).

Since the simulator RHS is real-valued, the complex tangent action is built from two real JVPs:

J Re[qhat exp(i kx x)] + i J Im[qhat exp(i kx x)].

Unlike a direct jacfwd over the real-packed vector [Re(qhat), Im(qhat)], this function assembles the complex Jacobian column-by-column in small batches. This avoids materializing all tangent directions at once.