Hamiltonian Mechanics in Polar Coordinates

import numpy as np 

def cartesian_to_polar(x, y):
    """
    Transform Cartesian coordinates to polar coordinates.
        :param x: Cartesian x-coordinate.
        :param y: Cartesian y-coordinate.
    :return: 
        A tuple (r, theta) where r is the radial coordinate,
        and theta is the angular coordinate.
    """
    r = np.sqrt(x**2 + y**2)
    theta = np.arctan2(y, x)
    return r, theta

def polar_velocity(r, theta, r_dot, theta_dot):
    """
    Calculate velocity components in polar coordinates.
        :param r: Radial coordinate.
        :param theta: Angular coordinate.
        :param r_dot: Derivative of radial coordinate.
        :param theta_dot: Derivative of angular coordinate.
    :return: 
        A tuple (v_r, v_theta) of velocity components.
    """
    v_r = r_dot
    v_theta = r * theta_dot
    return v_r, v_theta

def lagrangian_polar(m, r, theta, r_dot, theta_dot, V):
    """
        Compute the Lagrangian in polar coordinates.
        :param m: Mass of the object.
        :param r: Radial coordinate.
        :param theta: Angular coordinate.
        :param r_dot: Radial velocity.
        :param theta_dot: Angular velocity.
        :param V: Potential energy function V(r, theta).
    :return: 
        Lagrangian value.
    """
    T = 0.5 * m * (r_dot**2 + r**2 * theta_dot**2)
    return T - V(r, theta)

def hamiltonian_polar(m, p_r, p_theta, r, V):
    """
    Calculate the Hamiltonian in polar coordinates.
        :param m: Mass of the object.
        :param p_r: Radial conjugate momentum.
        :param p_theta: Angular conjugate momentum.
        :param r: Radial coordinate.
        :param V: Potential energy function V(r, theta).
    :return: 
        Hamiltonian value.
    """
    return p_r**2 / (2 * m) + p_theta**2 / (2 * m * r**2) + V(r) 

def dV_dr(r):
    return -1 / r**2
    
def update_polar_motion(r, theta, p_r, p_theta, m, V, dt):
    """
    Update the motion of a particle in polar coordinates
    numerical ly.
        :param r: Radial coordinate.
        :param theta: Angular coordinate.
        :param p_r: Radial momentum.
        :param p_theta: Angular momentum.
        :param m: Mass of the particle.
        :param V: Potential function as a function of r.
        :param dt: Time step for the update.
    :return: 
        Updated (r, theta, p_r, p_theta)
    """
# Compute derivatives from Hamilton's equations
    r_dot = p_r / m
    theta_dot = p_theta / (m * r**2)
    p_r_dot = p_theta**2 / (m * r**3) - dV_dr(r) # Example derived dV/dr
    p_theta_dot = 0 # Since V = V(r), partial V/partial theta = 0
    # Update positions and momenta
    r += r_dot * dt
    theta += theta_dot * dt
    p_r += p_r_dot * dt
    p_theta += p_theta_dot * dt
    return r, theta, p_r, p_theta

# Example of usage
m = 1.0
initial_r = 1.0
initial_theta = 0.0
initial_p_r = 1.0
initial_p_theta = 1.0

# Potential function example
V = lambda r: -1 * (1/r)
# Time step
dt = 0.01
# Update the motion
r, theta, p_r, p_theta = update_polar_motion(initial_r, initial_theta, initial_p_r, initial_p_theta, m, V, dt) 
print(f"Updated r: {r}, theta: (theta , p_r: {p_r}, p_theta: {p_theta}")

Updated r: 1.01, theta: (theta , p_r: 1.02, p_theta: 1.0