src
===

.. py:module:: src


Submodules
----------

.. toctree::
   :maxdepth: 1

   /autoapi/src/Refactoring/index
   /autoapi/src/equations/index
   /autoapi/src/plot/index
   /autoapi/src/solver/index


Functions
---------

.. autoapisummary::

   src.compute_dfm_dt
   src.objective
   src.solver_ode
   src.plot_results


Package Contents
----------------

.. py:function:: 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.

   :param f: jnp.ndarray
       The wave function represented as an array in the Fourier or physical space.
   :param k: float
       The wave number that determines the spatial frequency of the wave.
   :param nu: float
       The viscosity coefficient impacting energy dissipation.
   :param v_e: jnp.ndarray
       An array specifying velocity factors for the external forces.
   :param Nn: int
       The maximum number of wave modes in the system.
   :param index: int
       The current index in the array representing the mode being evaluated.
   :param q: jnp.ndarray
       An array of coefficients that modify the coupling between modes.

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


.. py:function:: 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.

   :param t: Current time, treated as a placeholder in this function.
   :type t: float
   :param f: Array of Fourier modes at the current timestep.
              Represents the state or amplitude values of the modes.
   :type f: jax.numpy.ndarray
   :param args: 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.
   :type args: tuple
   :return: Array containing the time derivatives of the Fourier modes computed over
            the range of indices specified by `m_max`.
   :rtype: jax.numpy.ndarray


.. py:function:: 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.

   :param k: Parameter to be passed as an argument to the ODE system.
   :type k: Array-like
   :param t_span: Time span over which the ODEs are solved. It must be a sequence
       containing the initial and final times of integration.
   :type t_span: Sequence[float]
   :param m_max: The maximum length used to initialize the state vector for the ODE
       problem.
   :type m_max: int
   :param v_e: Additional argument to be passed into the objective function.
   :type v_e: float
   :param nu: Additional argument to be passed into the objective function.
   :type nu: float
   :param q: Additional argument to be passed into the objective function.
   :type q: Array-like
   :param objective: The objective function defining the ODE system. It represents
       the time derivative of the state vector.
   :type objective: Callable
   :return: The solution of the ODE system, which includes the time-evaluated state
       vector at specified save points.
   :rtype: Array-like


.. py:function:: 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.

   :param k_values: A list or array of values for the variable k.
   :param slopes: A list or array of slope values corresponding to the k_values.
   :param 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.
   :param m_max: A numeric value representing the maximum m value used in the
       calculation, which will be displayed in the plot title.
   :return: None