Update stuff_model/calculate_arm_angles.py

stuff_model/calculate_arm_angles.py

@@ -36,6 +36,7 @@ def calculate_arm_angles(df: pl.DataFrame,pitcher_id:int) -> pl.DataFrame:
     else:
         shoulder_x = df_arm_angle.filter(pl.col("pitcher") == pitcher_id)["relative_shoulder_x"][0]
         shoulder_z = df_arm_angle.filter(pl.col("pitcher") == pitcher_id)["shoulder_z"][0]
+        ball_angle = df_arm_angle.filter(pl.col("pitcher") == pitcher_id)["ball_angle"][0]
 
         print(shoulder_x, shoulder_z)
 
@@ -55,6 +56,10 @@ def calculate_arm_angles(df: pl.DataFrame,pitcher_id:int) -> pl.DataFrame:
             #.drop(["Opp", "arm_angle_rad"])
             )
     
+        df_filter = df_filter.with_columns(
+                ((pl.col("arm_angle") * 0.25) + (ball_angle * 0.75)).alias("arm_angle")
+            )
+