astronomix.analysis_helpers.jacobians

`astronomix.analysis_helpers.jacobians`#

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.

Module Contents#

Functions#

`single_xmode_rhs_jacobian2D`	Return the dense +kx Fourier-block Jacobian of the simulator RHS.
`single_xmode_jacobian2Dt`	Return the dense +kx Fourier-block Jacobian of the time-integrated simulator.

API#

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.

Args:

primitive_state_unperturbed:: Primitive base state. It is converted to conserved variables before linearization.
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 still tight.

Returns:

J_complex:: Complex array of shape (nvar * Ny, nvar * Ny), representing the raw +kx Fourier-block tangent operator of the implemented simulator RHS.

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.