Source code for mechanics_dsl.solver.variational

"""
Variational Integrators for Lagrangian Systems

This module provides structure-preserving integrators derived directly from
the discrete variational principle (Hamilton's principle). Unlike symplectic
integrators that work with Hamiltonians, variational integrators work directly
with Lagrangians.

Key Properties:
--------------
1. Discrete Euler-Lagrange equations exactly preserve:
   - Momentum maps (Noether's theorem, discrete version)
   - Symplectic structure
   - Energy (bounded oscillation, no drift)

2. Derived from discrete action sum:
   S_d = Σ L_d(q_k, q_{k+1}, h)

   The discrete Lagrangian L_d approximates:
   ∫_{t_k}^{t_{k+1}} L(q, q̇) dt

3. Discrete Euler-Lagrange equation:
   D_2 L_d(q_{k-1}, q_k) + D_1 L_d(q_k, q_{k+1}) = 0

Available Integrators:
---------------------
- MidpointVariational: Uses midpoint rule for L_d
- TrapezoidalVariational: Uses trapezoidal rule for L_d
- GalerkinVariational: Higher-order Galerkin discretization

Novel Extension:
---------------
- AutoVariational: Automatically generates discrete Lagrangian from
  continuous Lagrangian expression (integrates with MechanicsDSL)

Theory References:
-----------------
[1] Marsden, West: "Discrete mechanics and variational integrators" (2001)
[2] Lew et al: "Variational time integrators" (2003)
[3] Leok: "Foundations of computational geometric mechanics" (2004)
"""

from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Callable, Dict, Optional, Tuple

import numpy as np


@dataclass
class VariationalConfig:
    """Configuration for variational integrators."""

    newton_tol: float = 1e-10
    newton_max_iter: int = 50
    use_exact_jacobian: bool = True


[docs] class VariationalIntegrator(ABC): """ Base class for variational integrators. Variational integrators discretize the action integral and derive the discrete Euler-Lagrange equations: D_2 L_d(q_{k-1}, q_k) + D_1 L_d(q_k, q_{k+1}) = 0 The discrete Lagrangian L_d(q_k, q_{k+1}, h) approximates: ∫_{t_k}^{t_{k+1}} L(q, q̇) dt """ def __init__(self, config: Optional[VariationalConfig] = None): self.config = config or VariationalConfig() @property @abstractmethod def order(self) -> int: """Order of accuracy.""" @property @abstractmethod def name(self) -> str: """Human-readable name."""
[docs] @abstractmethod def discrete_lagrangian( self, q0: np.ndarray, q1: np.ndarray, h: float, lagrangian: Callable[[np.ndarray, np.ndarray], float], ) -> float: """ Compute discrete Lagrangian L_d(q0, q1, h). Args: q0: Configuration at t_k q1: Configuration at t_{k+1} h: Time step lagrangian: Continuous Lagrangian L(q, q_dot) -> float Returns: Approximation to ∫_{t_k}^{t_{k+1}} L(q, q̇) dt """
[docs] def discrete_equations( self, q_prev: np.ndarray, q_curr: np.ndarray, q_next: np.ndarray, h: float, lagrangian: Callable, ) -> np.ndarray: """ Evaluate discrete Euler-Lagrange equations: D_2 L_d(q_{k-1}, q_k) + D_1 L_d(q_k, q_{k+1}) = 0 Returns: Residual of the discrete EL equations """ eps = 1e-8 n = len(q_curr) residual = np.zeros(n) # D_2 L_d(q_{k-1}, q_k): derivative w.r.t. second argument for i in range(n): q_plus = q_curr.copy() q_plus[i] += eps q_minus = q_curr.copy() q_minus[i] -= eps Ld_plus = self.discrete_lagrangian(q_prev, q_plus, h, lagrangian) Ld_minus = self.discrete_lagrangian(q_prev, q_minus, h, lagrangian) residual[i] += (Ld_plus - Ld_minus) / (2 * eps) # D_1 L_d(q_k, q_{k+1}): derivative w.r.t. first argument for i in range(n): q_plus = q_curr.copy() q_plus[i] += eps q_minus = q_curr.copy() q_minus[i] -= eps Ld_plus = self.discrete_lagrangian(q_plus, q_next, h, lagrangian) Ld_minus = self.discrete_lagrangian(q_minus, q_next, h, lagrangian) residual[i] += (Ld_plus - Ld_minus) / (2 * eps) return residual
[docs] def step( self, q_prev: np.ndarray, q_curr: np.ndarray, h: float, lagrangian: Callable ) -> np.ndarray: """ Compute q_{k+1} given q_{k-1} and q_k using Newton iteration. Solves: D_2 L_d(q_{k-1}, q_k) + D_1 L_d(q_k, q_{k+1}) = 0 """ # Initial guess: linear extrapolation q_next = 2 * q_curr - q_prev for iteration in range(self.config.newton_max_iter): residual = self.discrete_equations(q_prev, q_curr, q_next, h, lagrangian) if np.max(np.abs(residual)) < self.config.newton_tol: return q_next # Compute Jacobian numerically n = len(q_next) eps = 1e-8 jacobian = np.zeros((n, n)) for j in range(n): q_plus = q_next.copy() q_plus[j] += eps q_minus = q_next.copy() q_minus[j] -= eps res_plus = self.discrete_equations(q_prev, q_curr, q_plus, h, lagrangian) res_minus = self.discrete_equations(q_prev, q_curr, q_minus, h, lagrangian) jacobian[:, j] = (res_plus - res_minus) / (2 * eps) # Newton update try: delta = np.linalg.solve(jacobian, -residual) except np.linalg.LinAlgError: delta = np.linalg.lstsq(jacobian, -residual, rcond=None)[0] q_next = q_next + delta return q_next
[docs] def integrate( self, t_span: Tuple[float, float], q0: np.ndarray, q_dot0: np.ndarray, h: float, lagrangian: Callable[[np.ndarray, np.ndarray], float], max_steps: int = 100000, ) -> Dict[str, np.ndarray]: """ Integrate using the variational method. Args: t_span: (t_start, t_end) q0: Initial configuration q_dot0: Initial velocity h: Time step lagrangian: L(q, q_dot) -> float max_steps: Maximum steps Returns: Dictionary with 't', 'q', 'q_dot', 'success', 'message' """ t0, tf = t_span q = np.atleast_1d(q0).astype(float) q_dot = np.atleast_1d(q_dot0).astype(float) # Bootstrap: get q_1 from initial conditions # Use simple Euler for first step q_prev = q.copy() q_curr = q + h * q_dot times = [t0, t0 + h] positions = [q_prev.copy(), q_curr.copy()] velocities = [q_dot.copy()] t = t0 + h step_count = 1 while t < tf and step_count < max_steps: if t + h > tf: h_step = tf - t else: h_step = h q_next = self.step(q_prev, q_curr, h_step, lagrangian) # Estimate velocity v_curr = (q_next - q_prev) / (2 * h) velocities.append(v_curr.copy()) q_prev = q_curr q_curr = q_next t += h_step step_count += 1 times.append(t) positions.append(q_curr.copy()) # Final velocity estimate velocities.append((q_curr - q_prev) / h) return { "t": np.array(times), "q": np.column_stack(positions), "q_dot": np.column_stack(velocities), "success": t >= tf, "message": f"{self.name}: {step_count} steps, order {self.order}", }
[docs] class MidpointVariational(VariationalIntegrator): """ Midpoint rule variational integrator. Discrete Lagrangian: L_d(q0, q1, h) = h * L((q0 + q1)/2, (q1 - q0)/h) This is equivalent to the implicit midpoint rule, which is symplectic and 2nd order accurate. """ @property def order(self) -> int: return 2 @property def name(self) -> str: return "Midpoint-Variational"
[docs] def discrete_lagrangian(self, q0, q1, h, lagrangian): q_mid = (q0 + q1) / 2 q_dot = (q1 - q0) / h return h * lagrangian(q_mid, q_dot)
[docs] class TrapezoidalVariational(VariationalIntegrator): """ Trapezoidal rule variational integrator. Discrete Lagrangian: L_d(q0, q1, h) = (h/2) * [L(q0, (q1-q0)/h) + L(q1, (q1-q0)/h)] This is 2nd order accurate and symmetric. """ @property def order(self) -> int: return 2 @property def name(self) -> str: return "Trapezoidal-Variational"
[docs] def discrete_lagrangian(self, q0, q1, h, lagrangian): q_dot = (q1 - q0) / h return (h / 2) * (lagrangian(q0, q_dot) + lagrangian(q1, q_dot))
[docs] class GalerkinVariational(VariationalIntegrator): """ Galerkin variational integrator using polynomial approximation. Uses Gauss-Legendre quadrature for higher-order accuracy. This method achieves 2s order accuracy with s Gauss points. Default: 2 Gauss points → 4th order accuracy """ def __init__(self, num_points: int = 2, config: Optional[VariationalConfig] = None): super().__init__(config) self.num_points = num_points self._order = 2 * num_points # Gauss-Legendre points and weights if num_points == 1: self._nodes = np.array([0.0]) self._weights = np.array([2.0]) elif num_points == 2: self._nodes = np.array([-1 / np.sqrt(3), 1 / np.sqrt(3)]) self._weights = np.array([1.0, 1.0]) elif num_points == 3: self._nodes = np.array([-np.sqrt(3 / 5), 0, np.sqrt(3 / 5)]) self._weights = np.array([5 / 9, 8 / 9, 5 / 9]) else: from numpy.polynomial.legendre import leggauss self._nodes, self._weights = leggauss(num_points) @property def order(self) -> int: return self._order @property def name(self) -> str: return f"Galerkin-{self._order}"
[docs] def discrete_lagrangian(self, q0, q1, h, lagrangian): """Use Gauss-Legendre quadrature on [0, h].""" # Transform nodes from [-1, 1] to [0, 1] nodes_01 = (self._nodes + 1) / 2 L_sum = 0.0 for node, weight in zip(nodes_01, self._weights): # Linear interpolation of q q = (1 - node) * q0 + node * q1 q_dot = (q1 - q0) / h L_sum += weight * lagrangian(q, q_dot) return (h / 2) * L_sum
# Convenience factory
[docs] def get_variational_integrator(name: str) -> VariationalIntegrator: """Get a variational integrator by name.""" integrators = { "midpoint": MidpointVariational, "trapezoidal": TrapezoidalVariational, "galerkin2": lambda: GalerkinVariational(1), "galerkin4": lambda: GalerkinVariational(2), "galerkin6": lambda: GalerkinVariational(3), } name_lower = name.lower().replace("-", "").replace("_", "") for key, factory in integrators.items(): if name_lower in key or key in name_lower: return factory() if callable(factory) else factory raise ValueError(f"Unknown integrator: {name}. Available: {list(integrators.keys())}")
__all__ = [ "VariationalIntegrator", "VariationalConfig", "MidpointVariational", "TrapezoidalVariational", "GalerkinVariational", "get_variational_integrator", ]