Source code for mechanics_dsl.analysis.energy

"""
Energy Analysis Tools

Provides energy calculation and conservation analysis for mechanical systems.
"""

from typing import Callable, Dict, List, Optional

import numpy as np

from ..utils import logger


[docs] class EnergyAnalyzer: """ Energy analysis for mechanical systems. Computes kinetic energy, potential energy, and validates conservation. """ def __init__(self): self._kinetic_func: Optional[Callable] = None self._potential_func: Optional[Callable] = None
[docs] def compute_kinetic_energy( self, solution: dict, masses: Dict[str, float], velocities: Optional[List[str]] = None ) -> np.ndarray: """ Compute kinetic energy T = 0.5 * m * v^2 for each timestep. Args: solution: Simulation result masses: Dict mapping coordinate names to masses velocities: List of velocity variable names Returns: Kinetic energy array """ y = solution["y"] coords = solution.get("coordinates", []) n_points = len(solution["t"]) T = np.zeros(n_points) for i, q in enumerate(coords): m = masses.get(q, masses.get("m", 1.0)) v = y[2 * i + 1] # velocity is odd indices T += 0.5 * m * v**2 return T
[docs] def compute_potential_energy(self, solution: dict, potential_func: Callable) -> np.ndarray: """ Compute potential energy at each timestep. Args: solution: Simulation result potential_func: Function V(state) -> float Returns: Potential energy array """ y = solution["y"] n_points = y.shape[1] V = np.array([potential_func(y[:, i]) for i in range(n_points)]) return V
[docs] def check_conservation( self, solution: dict, kinetic: np.ndarray, potential: np.ndarray, tolerance: float = 1e-3 ) -> Dict: """ Check energy conservation. Args: solution: Simulation result kinetic: Kinetic energy array potential: Potential energy array tolerance: Relative tolerance for conservation check Returns: Dict with conservation analysis results """ total = kinetic + potential E0 = total[0] if abs(E0) < 1e-10: relative_error = np.abs(total - E0) else: relative_error = np.abs((total - E0) / E0) max_error = np.max(relative_error) mean_error = np.mean(relative_error) conserved = max_error < tolerance result = { "conserved": conserved, "initial_energy": E0, "max_relative_error": max_error, "mean_relative_error": mean_error, "total_energy": total, "kinetic_energy": kinetic, "potential_energy": potential, } if conserved: logger.info(f"Energy conserved: max error = {max_error:.2e}") else: logger.warning(f"Energy NOT conserved: max error = {max_error:.2e}") return result
[docs] def compute_pendulum_energy(self, solution: dict, m: float, length: float, g: float) -> Dict: """ Compute energy for simple pendulum. Args: solution: Simulation result m: Mass length: Pendulum length g: Gravitational acceleration Returns: Dict with kinetic, potential, and total energy """ y = solution["y"] theta = y[0] theta_dot = y[1] T = 0.5 * m * (length * theta_dot) ** 2 V = m * g * length * (1 - np.cos(theta)) return {"kinetic": T, "potential": V, "total": T + V}