src

Submodules

Functions

compute_dfm_dt(f, k, nu, v_e, Nn, index, q)

Computes the time derivative of a function representing a system described in the

objective(t, f, args)

Computes the time derivative of the Fourier mode dynamics based on the provided input

solver_ode(k, t_span, m_max, v_e, nu, q, objective)

Solves a system of ordinary differential equations (ODEs) using the specified

plot_results(k_values, slopes, mathematica_data, m_max)

Plots the results of the slope versus k value calculation alongside its

Package Contents

src.compute_dfm_dt(f, k, nu, v_e, Nn, index, q)

Computes the time derivative of a function representing a system described in the Fourier and physical space, taking into account effects like wave interactions, damping, and external forces.

Parameters:

  • f – jnp.ndarray The wave function represented as an array in the Fourier or physical space.

  • k – float The wave number that determines the spatial frequency of the wave.

  • nu – float The viscosity coefficient impacting energy dissipation.

  • v_e – jnp.ndarray An array specifying velocity factors for the external forces.

  • Nn – int The maximum number of wave modes in the system.

  • index – int The current index in the array representing the mode being evaluated.

  • q – jnp.ndarray An array of coefficients that modify the coupling between modes.

Returns:

jnp.complex128 The computed time derivative of the specified wave mode, taking into account interactions between different modes and external forces.

src.objective(t, f, args)

Computes the time derivative of the Fourier mode dynamics based on the provided input parameters and a set of arguments. This function utilizes jax.vmap for vectorized map computation, applying the compute_dfm_dt function over a specific range of modes.

The computation involves iterating over a sequence of values determined by the provided maximum number of modes (m_max) and processing elements using the arguments supplied in args. JAX’s vmap is employed to enable efficient parallel application of the calculation across multiple data points or modes.

The function encapsulates a physics-oriented calculation typically used in time-evolution simulations, spectrum modeling, or Fourier transform-based computational frameworks.

Parameters:

  • t (float) – Current time, treated as a placeholder in this function.

  • f (jax.numpy.ndarray) – Array of Fourier modes at the current timestep. Represents the state or amplitude values of the modes.

  • args (tuple) – Tuple of additional parameters passed to define the dynamics: - k: Wavenumber or characteristic parameter affecting dynamics. - v_e: Constant or coefficient influencing velocity evolution. - nu: Diffusion or damping coefficient used in the computation. - m_max: Maximum number of Fourier modes utilized in the summation. - q: External or forcing parameter influencing mode amplitudes.

Returns:

Array containing the time derivatives of the Fourier modes computed over the range of indices specified by m_max.

Return type:

jax.numpy.ndarray

src.solver_ode(k, t_span, m_max, v_e, nu, q, objective)

Solves a system of ordinary differential equations (ODEs) using the specified objective function and problem parameters. The solver integrates the system from a given initial state over a specified time span. It employs a Dopri5 solver with a PID-based step size controller to achieve precise numerical results.

Parameters:

  • k (Array-like) – Parameter to be passed as an argument to the ODE system.

  • t_span (Sequence[float]) – Time span over which the ODEs are solved. It must be a sequence containing the initial and final times of integration.

  • m_max (int) – The maximum length used to initialize the state vector for the ODE problem.

  • v_e (float) – Additional argument to be passed into the objective function.

  • nu (float) – Additional argument to be passed into the objective function.

  • q (Array-like) – Additional argument to be passed into the objective function.

  • objective (Callable) – The objective function defining the ODE system. It represents the time derivative of the state vector.

Returns:

The solution of the ODE system, which includes the time-evaluated state vector at specified save points.

Return type:

Array-like

src.plot_results(k_values, slopes, mathematica_data, m_max)

Plots the results of the slope versus k value calculation alongside its comparison with Mathematica data. The generated plot visualizes the relationship between the k values and the corresponding slopes, and it includes a display of Mathematica’s output for validation purposes.

Parameters:

  • k_values – A list or array of values for the variable k.

  • slopes – A list or array of slope values corresponding to the k_values.

  • mathematica_data – A DataFrame containing Mathematica’s results for slope vs. k, where the first column represents k values and the second column represents the associated slopes.

  • m_max – A numeric value representing the maximum m value used in the calculation, which will be displayed in the plot title.

Returns:

None