putter_modeler.py

# putter_modeler.py
# Calculates the physical properties of a putter from geometric components.
# Implements the logic for Phase 1 of the research plan.

import numpy as np

class Putter:
    """
    Represents a putter, calculating its mass, CG, and inertia tensor
    from a list of component parts.
    """
    def __init__(self, name, definition):
        """
        Initializes the Putter object.

        Args:
            name (str): The name of the putter (e.g., "Blade Putter").
            definition (dict): A dictionary from config.py defining the putter.
        """
        self.name = name
        self.components = definition["components"]
        self.shaft_attachment_point = definition["shaft_attachment_point"]
        
        self.total_mass = self._calculate_total_mass()
        self.center_of_gravity = self._calculate_cg()
        self.inertia_tensor = self._calculate_inertia_tensor()
        self.moi_around_shaft = self._calculate_moi_around_shaft()

    def _calculate_total_mass(self):
        """Calculates the total mass of the putter head."""
        return sum(c['mass'] for c in self.components)

    def _calculate_cg(self):
        """Calculates the center of gravity of the putter head."""
        weighted_sum = np.zeros(3)
        for c in self.components:
            weighted_sum += c['mass'] * c['pos']
        return weighted_sum / self.total_mass

    def _calculate_inertia_tensor(self):
        """
        Calculates the 3x3 inertia tensor of the composite body about its CG.
        Uses the Parallel Axis Theorem: I = I_cm + M(d^2 * I_unit - d x d)
        """
        I_total = np.zeros((3, 3))
        for c in self.components:
            m = c['mass']
            pos = c['pos']
            
            # 1. Inertia tensor of the component about its own center
            I_comp_cm = self._get_component_inertia(c)
            
            # 2. Use Parallel Axis Theorem to shift to the putter's total CG
            r = pos - self.center_of_gravity
            r_x, r_y, r_z = r
            
            # I_parallel_axis = M * [[y^2+z^2, -xy, -xz], [-yx, x^2+z^2, -yz], [-zx, -zy, x^2+y^2]]
            I_shift = m * (np.dot(r, r) * np.identity(3) - np.outer(r, r))
            
            I_total += I_comp_cm + I_shift
            
        return I_total

    def _get_component_inertia(self, component):
        """Calculates the inertia tensor for a single geometric component about its own CM."""
        m = component['mass']
        dims = component['dims']
        
        if component['shape'] == 'cuboid':
            l, w, h = dims # length, width, height
            Ixx = (1/12.0) * m * (w**2 + h**2)
            Iyy = (1/12.0) * m * (l**2 + h**2)
            Izz = (1/12.0) * m * (l**2 + w**2)
            return np.diag([Ixx, Iyy, Izz])
            
        elif component['shape'] == 'cylinder':
            # Assuming shaft axis is parallel to component's height axis (z)
            radius, height = dims
            Ixx = (1/12.0) * m * (3 * radius**2 + height**2)
            Iyy = (1/12.0) * m * (3 * radius**2 + height**2)
            Izz = 0.5 * m * radius**2 # MOI about cylinder axis
            return np.diag([Ixx, Iyy, Izz])
        
        return np.zeros((3,3))

    def _calculate_moi_around_shaft(self):
        """
        Calculates the effective Moment of Inertia around the shaft axis.
        This is the value that resists twisting during the stroke.
        Formula: I_shaft = u.T * I * u
        where u is the unit vector along the shaft axis.
        """
        # Assume a standard lie angle of 70 degrees, shaft points up and back
        lie_angle_rad = np.deg2rad(70)
        # Shaft vector in putter's coordinate system (assuming face is xy-plane)
        shaft_vector = np.array([0, -np.cos(lie_angle_rad), np.sin(lie_angle_rad)])
        shaft_unit_vector = shaft_vector / np.linalg.norm(shaft_vector)
        
        # We need the inertia tensor about the shaft attachment point, not the CG
        r = self.shaft_attachment_point - self.center_of_gravity
        I_shift = self.total_mass * (np.dot(r, r) * np.identity(3) - np.outer(r, r))
        I_at_shaft_point = self.inertia_tensor + I_shift

        # Project the inertia tensor onto the shaft axis
        moi = shaft_unit_vector.T @ I_at_shaft_point @ shaft_unit_vector
        return moi

    def __str__(self):
        """String representation for printing."""
        return (
            f"--- Putter Report: {self.name} ---\n"
            f"Total Mass: {self.total_mass:.3f} kg\n"
            f"Center of Gravity (CG): {np.round(self.center_of_gravity, 4)} m\n"
            f"MOI around shaft axis: {self.moi_around_shaft:.6f} kg*m^2\n"
            f"Inertia Tensor (about CG):\n{np.round(self.inertia_tensor, 6)}\n"
        )

if __name__ == '__main__':
    # This block runs if you execute this script directly for testing.
    from config import PUTTER_DEFINITIONS

    for name, definition in PUTTER_DEFINITIONS.items():
        putter = Putter(name, definition)
        print(putter)