leap_c.diff_mpc.data
====================

.. py:module:: leap_c.diff_mpc.data


Classes
-------

.. autoapisummary::

   leap_c.diff_mpc.data.AcadosOcpSolverInput


Functions
---------

.. autoapisummary::

   leap_c.diff_mpc.data.collate_acados_flattened_batch_iterate_fn
   leap_c.diff_mpc.data.collate_acados_flattened_iterate_fn
   leap_c.diff_mpc.data.collate_acados_ocp_solver_input
   leap_c.diff_mpc.data.validate_forward_inputs


Module Contents
---------------

.. py:class:: AcadosOcpSolverInput

   Bases: :py:obj:`NamedTuple`


   Input for an AcadosOcpSolver representing a batch of problem instances to be solved.

   :ivar x0: Initial state, shape (batch_size, nx)
   :ivar u0: Initial control input, shape (batch_size, nu), optional.
             If provided, the initial control input will be constrained to this.
   :ivar p_global: Global parameters, shape (batch_size, np_global), optional.
                   If not provided, the default values set in the acados ocp object will be used.
   :ivar p_stagewise: Stage-wise parameters, shape (batch_size, N_horizon + 1, np_stagewise),
                      or (batch_size, len(p_stagewise_sparse_idx), np_stagewise), if p_stagewise_sparse_idx is
                      provided, optional.
                      If not provided, the default values set in the acados ocp object will be used.
                      Has to be provided if p_stagewise_sparse_idx is provided.
   :ivar p_stagewise_sparse_idx: If provided, the indices determine which elements of the
                                 total stagewise parameter in the solver should be overwritten
                                 by the provided p_stagewise values,
                                 shape (batch_size, N_horizon + 1, nindices), optional.



   .. py:method:: get_sample(idx: int) -> AcadosOcpSolverInput

      Get the sample at index `idx` from the batch.



   .. py:property:: batch_size
      :type: int


      Get the batch size.


   .. py:attribute:: p_global
      :type:  numpy.ndarray | None
      :value: None



   .. py:attribute:: p_stagewise
      :type:  numpy.ndarray | None
      :value: None



   .. py:attribute:: p_stagewise_sparse_idx
      :type:  numpy.ndarray | None
      :value: None



   .. py:attribute:: u0
      :type:  numpy.ndarray | None
      :value: None



   .. py:attribute:: x0
      :type:  numpy.ndarray


.. py:function:: collate_acados_flattened_batch_iterate_fn(batch: Sequence[acados_template.acados_ocp_iterate.AcadosOcpFlattenedBatchIterate], collate_fn_map: dict | None = None) -> acados_template.acados_ocp_iterate.AcadosOcpFlattenedBatchIterate

.. py:function:: collate_acados_flattened_iterate_fn(batch: Sequence[acados_template.acados_ocp_iterate.AcadosOcpFlattenedIterate], collate_fn_map: dict | None = None) -> acados_template.acados_ocp_iterate.AcadosOcpFlattenedBatchIterate

.. py:function:: collate_acados_ocp_solver_input(batch: Sequence[AcadosOcpSolverInput], collate_fn_map: dict | None = None) -> AcadosOcpSolverInput

   Collates a batch of AcadosOcpSolverInput objects into a single object.


.. py:function:: validate_forward_inputs(ocp: acados_template.AcadosOcp, x0: numpy.ndarray | torch.Tensor, u0: numpy.ndarray | torch.Tensor | None) -> None

   Validate the shapes of ``x0`` and ``u0`` before the forward solve.

   Works on any array-like exposing ``ndim`` and ``shape`` (numpy, torch, jax), so it can be
   shared by the torch and jax layers. Run before solver initialization so dimension mismatches
   fail fast with a clear message rather than as a cryptic C-level error.

   :raises ValueError: If ``x0`` is not 2-D, its last dimension is not ``ocp.dims.nx``, ``u0`` is
       not 2-D, its batch does not match ``x0``, or its last dimension is not ``ocp.dims.nu``.