simulation_engine.py

# simulation_engine.py
# Core logic for simulating a single golf putt, including stroke, noise, and impact.
# Implements the logic for Phase 2 of the research plan.

import numpy as np
import config

def simulate_one_putt(putter, setup_noise, path_noise_func, face_noise_func):
    """
    Simulates a single putting stroke and the subsequent ball impact.

    Args:
        putter (Putter): The putter object to use.
        setup_noise (dict): A dictionary containing noise values for this specific putt.
        path_noise_func (function): A function that returns random path noise.
        face_noise_func (function): A function that returns random face torque noise.

    Returns:
        tuple: (ball_initial_velocity, ball_initial_spin) vectors.
    """
    
    # --- 1. STROKE SIMULATION ---
    # Model the stroke as a simple pendulum to get impact velocity.
    # A more complex ODE model could be swapped in here.
    pivot = config.SIMULATION_SETUP["swing_pivot_point"]
    g = config.G
    
    # Distance from pivot to putter's CG
    r_pivot_to_cg = np.linalg.norm(putter.center_of_gravity - pivot)
    
    # Use conservation of energy for a simple pendulum to find max speed
    # mgh = 0.5 * I * w^2 => w = sqrt(2mgh / I)
    # where h is the vertical drop from the top of the swing.
    theta_max = np.deg2rad(config.SIMULATION_SETUP["swing_amplitude_deg"])
    h = r_pivot_to_cg * (1 - np.cos(theta_max))
    
    # For a simple pendulum, omega = sqrt(2gh) / r
    # This gives us the angular velocity at the bottom of the swing.
    omega_impact = np.sqrt(2 * g * h) / r_pivot_to_cg
    
    # Putter head's linear velocity at impact (v = r * w)
    # Assuming the swing is in the y-z plane (straight back, straight through)
    putter_vel_at_impact = np.array([0, r_pivot_to_cg * omega_impact, 0])
    
    # --- 2. APPLY BIOMECHANICAL NOISE ---
    # The noise functions generate random values for this specific putt trial
    path_force = path_noise_func()
    face_torque = face_noise_func()
    
    # Apply noise to determine the exact impact location and face angle
    # A simple model: torque noise creates an angular deviation at impact
    # face_torque = I * alpha => alpha = face_torque / I_shaft
    # delta_angle = 0.5 * alpha * t^2 (assuming constant torque over downswing)
    downswing_time_approx = 0.5 * (np.pi * np.sqrt(r_pivot_to_cg / g)) # half period
    alpha = face_torque / putter.moi_around_shaft
    face_angle_error_rad = 0.5 * alpha * downswing_time_approx**2 + np.deg2rad(setup_noise['setup_angle'])
    
    # Path noise can create a small off-center hit location
    # F = ma => a = F/m. delta_x = 0.5 * a * t^2
    acceleration_off_path = path_force / putter.total_mass
    impact_offset_x = 0.5 * acceleration_off_path * downswing_time_approx**2

    # --- 3. IMPACT SIMULATION ---
    # Now we model the 2D collision in the x-y plane of the putter face
    
    v_putter_pre = putter_vel_at_impact[1] # Speed into the ball
    m_putter = putter.total_mass
    m_ball = config.GOLF_BALL_MASS
    cor = config.COR
    
    # MOI for rotation around the vertical (z) axis through the CG
    I_putter_z = putter.inertia_tensor[2, 2]
    
    # The impact happens at x = impact_offset_x
    d = impact_offset_x
    
    # Solve the system of equations for collision:
    # 1. Linear Momentum: m_putter*v_putter_pre = m_putter*v_putter_post + m_ball*v_ball_post
    # 2. Angular Momentum: d*J = I*w_post (where J is impulse)
    # 3. Impulse on ball: J = m_ball*v_ball_post
    # 4. Relative velocity (Newton's Law of Restitution): v_ball_post - (v_putter_post + d*w_post) = -cor * v_putter_pre
    
    # Solving these gives the final velocity of the ball
    numerator = v_putter_pre * m_putter * (1 + cor)
    denominator = m_putter + m_ball + (m_putter * m_ball * d**2) / I_putter_z
    v_ball_post_y = numerator / denominator
    
    # The final velocity of the putter head's CG
    v_putter_post_y = (m_putter * v_putter_pre - m_ball * v_ball_post_y) / m_putter

    # The angular velocity of the putter after impact
    w_putter_post_z = (m_ball * v_ball_post_y * d) / I_putter_z

    # The ball also gets a velocity component due to the face rotating at impact
    # This is the "gear effect"
    v_ball_post_x = d * w_putter_post_z

    # Ball's final velocity vector, also rotated by the initial face angle error
    v_ball_unrotated = np.array([v_ball_post_x, v_ball_post_y, 0])
    
    c, s = np.cos(face_angle_error_rad), np.sin(face_angle_error_rad)
    rotation_matrix = np.array([[c, -s, 0], [s, c, 0], [0, 0, 1]])
    
    ball_initial_velocity = rotation_matrix @ v_ball_unrotated
    
    # For this model, we'll assume no significant spin is generated, but a more
    # complex model could calculate it here.
    ball_initial_spin = np.zeros(3)
    
    return ball_initial_velocity, ball_initial_spin