"""
Energy calculation utilities for MechanicsDSL
"""
from typing import Dict
import numpy as np
from .utils import logger, validate_solution_dict
[docs]
class PotentialEnergyCalculator:
"""Compute potential energy with proper offset for different systems"""
[docs]
@staticmethod
def compute_pe_offset(system_type: str, parameters: Dict[str, float]) -> float:
"""
Compute PE offset to set minimum PE = 0
Args:
system_type: Type of mechanical system
parameters: System parameters
Returns:
PE offset value
"""
system = system_type.lower()
if "double" in system and "pendulum" in system:
m1 = parameters.get("m1", 1.0)
m2 = parameters.get("m2", 1.0)
l1 = parameters.get("l1", 1.0)
l2 = parameters.get("l2", 1.0)
g = parameters.get("g", 9.81)
# Minimum PE when both pendulums hang straight down
return -m1 * g * l1 - m2 * g * (l1 + l2)
elif "spherical" in system and "pendulum" in system:
m = parameters.get("m", 1.0)
length = parameters.get("l", 1.0)
g = parameters.get("g", 9.81)
# Spherical pendulum: minimum PE when theta=0 (hanging straight down)
return 0.0 # PE = mgl(1-cos(theta)), already 0 at theta=0
elif "pendulum" in system:
m = parameters.get("m", 1.0)
length = parameters.get("l", 1.0)
g = parameters.get("g", 9.81)
# Minimum PE when pendulum hangs straight down
return -m * g * length
elif "oscillator" in system or "spring" in system:
# Harmonic oscillator: PE minimum is already at x=0
return 0.0
else:
# Default: no offset
return 0.0
[docs]
@staticmethod
def compute_kinetic_energy(solution: dict, parameters: Dict[str, float]) -> np.ndarray:
"""
Compute kinetic energy from solution with validation.
Args:
solution: Solution dictionary (validated)
parameters: System parameters dictionary
Returns:
Array of kinetic energy values
Raises:
TypeError: If inputs have wrong types
ValueError: If solution is invalid
"""
if not isinstance(parameters, dict):
raise TypeError(f"parameters must be dict, got {type(parameters).__name__}")
validate_solution_dict(solution)
t = solution["t"]
y = solution["y"]
coords = solution["coordinates"]
KE = np.zeros_like(t)
if len(coords) == 0:
logger.warning("No coordinates found for kinetic energy calculation")
return KE
if "theta" in coords[0]: # Pendulum systems
if len(coords) == 1: # Simple pendulum
if y.shape[0] < 2:
logger.warning("Insufficient state vector for simple pendulum KE")
return KE
theta_dot = y[1]
m = parameters.get("m", 1.0)
length = parameters.get("l", 1.0)
KE = 0.5 * m * length**2 * theta_dot**2
elif len(coords) >= 2: # Double or spherical pendulum
if y.shape[0] < 4:
logger.warning("Insufficient state vector for multi-coord pendulum KE")
return KE
# Check for spherical pendulum (theta, phi) vs double pendulum (theta1, theta2)
if "phi" in coords:
# Spherical pendulum: T = 0.5*m*l^2*(theta_dot^2 + sin^2(theta)*phi_dot^2)
theta = y[0]
theta_dot = y[1]
phi_dot = y[3]
m = parameters.get("m", 1.0)
length = parameters.get("l", 1.0)
KE = 0.5 * m * length**2 * (theta_dot**2 + np.sin(theta) ** 2 * phi_dot**2)
else:
# Double pendulum
theta1_dot, theta2_dot = y[1], y[3]
theta1, theta2 = y[0], y[2]
m1 = parameters.get("m1", 1.0)
m2 = parameters.get("m2", 1.0)
l1 = parameters.get("l1", 1.0)
l2 = parameters.get("l2", 1.0)
KE1 = 0.5 * m1 * l1**2 * theta1_dot**2
KE2 = (
0.5
* m2
* (
l1**2 * theta1_dot**2
+ l2**2 * theta2_dot**2
+ 2 * l1 * l2 * theta1_dot * theta2_dot * np.cos(theta1 - theta2)
)
)
KE = KE1 + KE2
else: # Cartesian systems
v = y[1] if y.shape[0] > 1 else np.zeros_like(t)
m = parameters.get("m", 1.0)
KE = 0.5 * m * v**2
return KE
[docs]
@staticmethod
def compute_potential_energy(
solution: dict, parameters: Dict[str, float], system_type: str = ""
) -> np.ndarray:
"""Compute potential energy from solution with proper offset"""
t = solution["t"]
y = solution["y"]
coords = solution["coordinates"]
PE = np.zeros_like(t)
if len(coords) == 0:
logger.warning("No coordinates found for potential energy calculation")
return PE
if "theta" in coords[0]: # Pendulum systems
if len(coords) == 1: # Simple pendulum
if y.shape[0] < 1:
logger.warning("Insufficient state vector for simple pendulum PE")
return PE
theta = y[0]
m = parameters.get("m", 1.0)
length = parameters.get("l", 1.0)
g = parameters.get("g", 9.81)
PE = -m * g * length * np.cos(theta)
offset = PotentialEnergyCalculator.compute_pe_offset("simple_pendulum", parameters)
PE = PE - offset
elif len(coords) >= 2: # Double or spherical pendulum
if y.shape[0] < 3:
logger.warning("Insufficient state vector for multi-coord pendulum PE")
return PE
# Check for spherical pendulum (theta, phi) vs double pendulum (theta1, theta2)
if "phi" in coords or "spherical" in system_type.lower():
# Spherical pendulum: V = mgl(1 - cos(theta))
theta = y[0]
m = parameters.get("m", 1.0)
length = parameters.get("l", 1.0)
g = parameters.get("g", 9.81)
PE = m * g * length * (1 - np.cos(theta))
else:
# Double pendulum
theta1, theta2 = y[0], y[2]
m1 = parameters.get("m1", 1.0)
m2 = parameters.get("m2", 1.0)
l1 = parameters.get("l1", 1.0)
l2 = parameters.get("l2", 1.0)
g = parameters.get("g", 9.81)
PE1 = -m1 * g * l1 * np.cos(theta1)
PE2 = -m2 * g * (l1 * np.cos(theta1) + l2 * np.cos(theta2))
PE = PE1 + PE2
offset = PotentialEnergyCalculator.compute_pe_offset(
"double_pendulum", parameters
)
PE = PE - offset
else: # Cartesian systems
if y.shape[0] < 1:
logger.warning("Insufficient state vector for Cartesian PE")
return PE
x = y[0]
k = parameters.get("k", 1.0)
PE = 0.5 * k * x**2
return PE