metrics.py

import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
import numpy as np
import tensorflow as tf
MAX_OUTPUT_FRAMES_V18 = 30
NFL_FIELD_X_MAX = 120.0
NFL_FIELD_Y_MAX = 53.3
def seq2seq_masked_rmse_loss(y_true, y_pred):
    y_true = tf.where(tf.math.is_finite(y_true), y_true, tf.zeros_like(y_true))
    y_pred = tf.where(tf.math.is_finite(y_pred), y_pred, tf.zeros_like(y_pred))
    y_true_rank = tf.rank(y_true)
    def process_y_true_90():
        y_true_90 = tf.reshape(y_true, [-1, 90])
        y_true_x = y_true_90[:, 0:MAX_OUTPUT_FRAMES_V18]
        y_true_y = y_true_90[:, MAX_OUTPUT_FRAMES_V18:MAX_OUTPUT_FRAMES_V18*2]
        y_true_coords = tf.stack([y_true_x, y_true_y], axis=2)
        y_mask = y_true_90[:, MAX_OUTPUT_FRAMES_V18*2:MAX_OUTPUT_FRAMES_V18*3]
        return y_true_coords, y_mask
    def process_y_true_30_2():
        y_true_coords = tf.reshape(y_true, [-1, MAX_OUTPUT_FRAMES_V18, 2])
        coords_finite = tf.cast(tf.math.is_finite(y_true_coords), tf.float32)
        coords_in_range = tf.cast(
            (y_true_coords[:, :, 0] >= 0.0) & (y_true_coords[:, :, 0] <= NFL_FIELD_X_MAX) &
            (y_true_coords[:, :, 1] >= 0.0) & (y_true_coords[:, :, 1] <= NFL_FIELD_Y_MAX),
            tf.float32
        )
        coords_in_range_expanded = tf.expand_dims(coords_in_range, axis=-1)
        coords_valid = coords_finite * coords_in_range_expanded
        y_mask = tf.reduce_max(coords_valid, axis=-1)
        return y_true_coords, y_mask
    y_true_coords, y_mask = tf.cond(
        tf.equal(y_true_rank, 3),
        process_y_true_30_2,
        process_y_true_90
    )
    y_mask = tf.clip_by_value(y_mask, 0.0, 1.0)
    y_pred_rank = tf.rank(y_pred)
    y_pred = tf.cond(
        tf.equal(y_pred_rank, 3),
        lambda: y_pred,
        lambda: tf.reshape(y_pred, [-1, MAX_OUTPUT_FRAMES_V18, 2])
    )
    diff = y_true_coords - y_pred
    diff = tf.clip_by_value(diff, -1e6, 1e6)
    squared_error = tf.square(diff)
    y_mask_expanded = tf.expand_dims(y_mask, axis=-1)
    masked_squared_error = squared_error * y_mask_expanded
    total_valid_frames = tf.reduce_sum(y_mask) * 2.0
    total_valid_frames = tf.maximum(total_valid_frames, 1.0)
    mse = tf.reduce_sum(masked_squared_error) / total_valid_frames
    y_true_directions = y_true_coords[:, 1:, :] - y_true_coords[:, :-1, :]
    y_pred_directions = y_pred[:, 1:, :] - y_pred[:, :-1, :]
    y_true_dir_norm = tf.sqrt(tf.reduce_sum(tf.square(y_true_directions), axis=-1, keepdims=True) + 1e-8)
    y_pred_dir_norm = tf.sqrt(tf.reduce_sum(tf.square(y_pred_directions), axis=-1, keepdims=True) + 1e-8)
    y_true_dir_normalized = y_true_directions / y_true_dir_norm
    y_pred_dir_normalized = y_pred_directions / y_pred_dir_norm
    direction_cosine = tf.reduce_sum(y_true_dir_normalized * y_pred_dir_normalized, axis=-1)
    direction_cosine = tf.clip_by_value(direction_cosine, -1.0, 1.0)
    direction_mask = y_mask[:, :-1]
    direction_loss = tf.reduce_sum((1.0 - direction_cosine) * direction_mask) / tf.maximum(tf.reduce_sum(direction_mask), 1.0)
    rmse = tf.sqrt(mse + 1e-8)
    rmse = tf.where(tf.math.is_finite(rmse), rmse, tf.constant(1.0, dtype=rmse.dtype))
    combined_loss = rmse + 0.1 * direction_loss
    return combined_loss
def create_physics_informed_loss(y_scaler, vel_weight=0.2, accel_weight=0.05, boundary_weight=0.1, start_weight=0.1):
    scale_ = y_scaler.scale_
    scale_x = np.mean(scale_[:30])
    scale_y = np.mean(scale_[30:])
    scale_avg = (scale_x + scale_y) / 2.0
    MAX_VEL_YARDS_PER_FRAME = 1.2
    MAX_VEL_SCALED = MAX_VEL_YARDS_PER_FRAME / scale_avg
    mean_ = y_scaler.mean_
    mean_y = np.mean(mean_[30:])
    scale_y_val = np.mean(scale_[30:])
    Y_MIN_YARDS = 0.0
    Y_MAX_YARDS = 53.3
    Y_MIN_SCALED = (Y_MIN_YARDS - mean_y) / scale_y_val
    Y_MAX_SCALED = (Y_MAX_YARDS - mean_y) / scale_y_val
    MAX_VEL_SCALED_TF = tf.constant(MAX_VEL_SCALED, dtype=tf.float32)
    Y_MIN_SCALED_TF = tf.constant(Y_MIN_SCALED, dtype=tf.float32)
    Y_MAX_SCALED_TF = tf.constant(Y_MAX_SCALED, dtype=tf.float32)
    def physics_informed_loss(y_true, y_pred):
        y_true = tf.where(tf.math.is_finite(y_true), y_true, tf.zeros_like(y_true))
        y_pred = tf.where(tf.math.is_finite(y_pred), y_pred, tf.zeros_like(y_pred))
        y_true_rank = tf.rank(y_true)
        def process_y_true_90():
            y_true_90 = tf.reshape(y_true, [-1, 90])
            y_true_x = y_true_90[:, 0:MAX_OUTPUT_FRAMES_V18]
            y_true_y = y_true_90[:, MAX_OUTPUT_FRAMES_V18:MAX_OUTPUT_FRAMES_V18*2]
            y_true_coords = tf.stack([y_true_x, y_true_y], axis=2)
            y_mask = y_true_90[:, MAX_OUTPUT_FRAMES_V18*2:MAX_OUTPUT_FRAMES_V18*3]
            return y_true_coords, y_mask
        def process_y_true_30_2():
            y_true_coords = tf.reshape(y_true, [-1, MAX_OUTPUT_FRAMES_V18, 2])
            coords_finite = tf.cast(tf.math.is_finite(y_true_coords), tf.float32)
            y_mask = tf.reduce_max(coords_finite, axis=-1)
            return y_true_coords, y_mask
        y_true_coords, y_mask = tf.cond(
            tf.equal(y_true_rank, 3),
            process_y_true_30_2,
            process_y_true_90
        )
        y_mask = tf.clip_by_value(y_mask, 0.0, 1.0)
        y_pred_rank = tf.rank(y_pred)
        y_pred_shaped = tf.cond(
            tf.equal(y_pred_rank, 3),
            lambda: y_pred,
            lambda: tf.reshape(y_pred, [-1, MAX_OUTPUT_FRAMES_V18, 2])
        )
        diff = y_true_coords - y_pred_shaped
        diff = tf.clip_by_value(diff, -1e6, 1e6)
        squared_error = tf.square(diff)
        y_mask_expanded = tf.expand_dims(y_mask, axis=-1)
        masked_squared_error = squared_error * y_mask_expanded
        total_valid_frames = tf.reduce_sum(y_mask) * 2.0
        total_valid_frames = tf.maximum(total_valid_frames, 1.0)
        mse = tf.reduce_sum(masked_squared_error) / total_valid_frames
        rmse = tf.sqrt(mse + 1e-8)
        rmse = tf.where(tf.math.is_finite(rmse), rmse, tf.constant(1.0, dtype=rmse.dtype))
        y_true_directions = y_true_coords[:, 1:, :] - y_true_coords[:, :-1, :]
        y_pred_directions = y_pred_shaped[:, 1:, :] - y_pred_shaped[:, :-1, :]
        y_true_dir_norm = tf.sqrt(tf.reduce_sum(tf.square(y_true_directions), axis=-1, keepdims=True) + 1e-8)
        y_pred_dir_norm = tf.sqrt(tf.reduce_sum(tf.square(y_pred_directions), axis=-1, keepdims=True) + 1e-8)
        y_true_dir_normalized = y_true_directions / y_true_dir_norm
        y_pred_dir_normalized = y_pred_directions / y_pred_dir_norm
        direction_cosine = tf.reduce_sum(y_true_dir_normalized * y_pred_dir_normalized, axis=-1)
        direction_cosine = tf.clip_by_value(direction_cosine, -1.0, 1.0)
        direction_mask = y_mask[:, :-1]
        direction_loss = tf.reduce_sum((1.0 - direction_cosine) * direction_mask) / tf.maximum(tf.reduce_sum(direction_mask), 1.0)
        velocity = y_pred_shaped[:, 1:, :] - y_pred_shaped[:, :-1, :]
        vel_mag = tf.sqrt(tf.reduce_sum(tf.square(velocity), axis=-1) + 1e-8)
        vel_penalty = tf.reduce_mean(tf.nn.relu(vel_mag - MAX_VEL_SCALED_TF))
        accel = y_pred_shaped[:, 2:, :] - 2.0 * y_pred_shaped[:, 1:-1, :] + y_pred_shaped[:, :-2, :]
        accel_penalty = tf.reduce_mean(tf.square(accel))
        y_coords = y_pred_shaped[:, :, 1]
        boundary_penalty = tf.reduce_mean(tf.nn.relu(Y_MIN_SCALED_TF - y_coords)) + \
                          tf.reduce_mean(tf.nn.relu(y_coords - Y_MAX_SCALED_TF))
        start_penalty = tf.reduce_mean(tf.square(y_pred_shaped[:, 0, :]))
        total_loss = rmse + \
                    0.1 * direction_loss + \
                    vel_weight * vel_penalty + \
                    accel_weight * accel_penalty + \
                    boundary_weight * boundary_penalty + \
                    start_weight * start_penalty
        return total_loss
    return physics_informed_loss
def create_physics_informed_loss_for_delta(y_scaler, vel_weight=0.1, accel_weight=0.05):
    scale_ = y_scaler.scale_
    scale_avg = np.mean(scale_)
    MAX_VEL_YARDS_PER_FRAME = 1.2
    MAX_VEL_SCALED = MAX_VEL_YARDS_PER_FRAME / scale_avg if scale_avg > 0 else 3.0
    MAX_VEL_SCALED_TF = tf.constant(MAX_VEL_SCALED, dtype=tf.float32)
    def physics_informed_loss_for_delta(y_true, y_pred):
        y_true = tf.where(tf.math.is_finite(y_true), y_true, tf.zeros_like(y_true))
        y_pred = tf.where(tf.math.is_finite(y_pred), y_pred, tf.zeros_like(y_pred))
        y_true_rank = tf.rank(y_true)
        def process_y_true_90():
            y_true_90 = tf.reshape(y_true, [-1, 90])
            y_true_x = y_true_90[:, 0:MAX_OUTPUT_FRAMES_V18]
            y_true_y = y_true_90[:, MAX_OUTPUT_FRAMES_V18:MAX_OUTPUT_FRAMES_V18*2]
            y_true_coords = tf.stack([y_true_x, y_true_y], axis=2)
            y_mask = y_true_90[:, MAX_OUTPUT_FRAMES_V18*2:MAX_OUTPUT_FRAMES_V18*3]
            return y_true_coords, y_mask
        def process_y_true_30_2():
            y_true_coords = tf.reshape(y_true, [-1, MAX_OUTPUT_FRAMES_V18, 2])
            coords_finite = tf.cast(tf.math.is_finite(y_true_coords), tf.float32)
            y_mask = tf.reduce_max(coords_finite, axis=-1)
            return y_true_coords, y_mask
        y_true_coords, y_mask = tf.cond(
            tf.equal(y_true_rank, 3),
            process_y_true_30_2,
            process_y_true_90
        )
        y_mask = tf.clip_by_value(y_mask, 0.0, 1.0)
        y_pred_rank = tf.rank(y_pred)
        y_pred_shaped = tf.cond(
            tf.equal(y_pred_rank, 3),
            lambda: y_pred,
            lambda: tf.reshape(y_pred, [-1, MAX_OUTPUT_FRAMES_V18, 2])
        )
        diff = y_true_coords - y_pred_shaped
        diff = tf.clip_by_value(diff, -1e6, 1e6)
        squared_error = tf.square(diff)
        y_mask_expanded = tf.expand_dims(y_mask, axis=-1)
        masked_squared_error = squared_error * y_mask_expanded
        total_valid_frames = tf.reduce_sum(y_mask) * 2.0
        total_valid_frames = tf.maximum(total_valid_frames, 1.0)
        mse = tf.reduce_sum(masked_squared_error) / total_valid_frames
        rmse = tf.sqrt(mse + 1e-8)
        rmse = tf.where(tf.math.is_finite(rmse), rmse, tf.constant(1.0, dtype=rmse.dtype))
        y_true_directions = y_true_coords[:, 1:, :]
        y_pred_directions = y_pred_shaped[:, 1:, :]
        y_true_dir_norm = tf.sqrt(tf.reduce_sum(tf.square(y_true_directions), axis=-1, keepdims=True) + 1e-8)
        y_pred_dir_norm = tf.sqrt(tf.reduce_sum(tf.square(y_pred_directions), axis=-1, keepdims=True) + 1e-8)
        y_true_dir_normalized = y_true_directions / y_true_dir_norm
        y_pred_dir_normalized = y_pred_directions / y_pred_dir_norm
        direction_cosine = tf.reduce_sum(y_true_dir_normalized * y_pred_dir_normalized, axis=-1)
        direction_cosine = tf.clip_by_value(direction_cosine, -1.0, 1.0)
        direction_mask = y_mask[:, 1:]
        direction_loss = tf.reduce_sum((1.0 - direction_cosine) * direction_mask) / tf.maximum(tf.reduce_sum(direction_mask), 1.0)
        velocity = y_pred_shaped[:, 1:, :]
        speed_mag = tf.sqrt(tf.reduce_sum(tf.square(velocity), axis=-1) + 1e-8)
        vel_penalty = tf.reduce_mean(tf.nn.relu(speed_mag - MAX_VEL_SCALED_TF))
        accel = velocity[:, 1:, :] - velocity[:, :-1, :]
        accel_penalty = tf.reduce_mean(tf.square(accel))
        start_penalty = tf.reduce_mean(tf.square(y_pred_shaped[:, 0, :]))
        total_loss = rmse + \
                    0.1 * direction_loss + \
                    vel_weight * vel_penalty + \
                    accel_weight * accel_penalty + \
                    0.05 * start_penalty
        return total_loss
    return physics_informed_loss_for_delta
def seq2seq_masked_rmse_metric(y_true, y_pred):
    return seq2seq_masked_rmse_loss(y_true, y_pred)
def seq2seq_pure_rmse_metric(y_true, y_pred):
    y_true = tf.where(tf.math.is_finite(y_true), y_true, tf.zeros_like(y_true))
    y_pred = tf.where(tf.math.is_finite(y_pred), y_pred, tf.zeros_like(y_pred))
    y_true_rank = tf.rank(y_true)
    def process_y_true_90():
        y_true_90 = tf.reshape(y_true, [-1, 90])
        y_true_x = y_true_90[:, 0:MAX_OUTPUT_FRAMES_V18]
        y_true_y = y_true_90[:, MAX_OUTPUT_FRAMES_V18:MAX_OUTPUT_FRAMES_V18*2]
        y_true_coords = tf.stack([y_true_x, y_true_y], axis=2)
        y_mask = y_true_90[:, MAX_OUTPUT_FRAMES_V18*2:MAX_OUTPUT_FRAMES_V18*3]
        return y_true_coords, y_mask
    def process_y_true_30_2():
        y_true_coords = tf.reshape(y_true, [-1, MAX_OUTPUT_FRAMES_V18, 2])
        coords_finite = tf.cast(tf.math.is_finite(y_true_coords), tf.float32)
        coords_in_range = tf.cast(
            (y_true_coords[:, :, 0] >= 0.0) & (y_true_coords[:, :, 0] <= NFL_FIELD_X_MAX) &
            (y_true_coords[:, :, 1] >= 0.0) & (y_true_coords[:, :, 1] <= NFL_FIELD_Y_MAX),
            tf.float32
        )
        coords_in_range_expanded = tf.expand_dims(coords_in_range, axis=-1)
        coords_valid = coords_finite * coords_in_range_expanded
        y_mask = tf.reduce_max(coords_valid, axis=-1)
        return y_true_coords, y_mask
    y_true_coords, y_mask = tf.cond(
        tf.equal(y_true_rank, 3),
        process_y_true_30_2,
        process_y_true_90
    )
    y_mask = tf.clip_by_value(y_mask, 0.0, 1.0)
    y_pred_rank = tf.rank(y_pred)
    y_pred = tf.cond(
        tf.equal(y_pred_rank, 3),
        lambda: y_pred,
        lambda: tf.reshape(y_pred, [-1, MAX_OUTPUT_FRAMES_V18, 2])
    )
    diff = y_true_coords - y_pred
    diff = tf.clip_by_value(diff, -1e6, 1e6)
    squared_error = tf.square(diff)
    y_mask_expanded = tf.expand_dims(y_mask, axis=-1)
    masked_squared_error = squared_error * y_mask_expanded
    total_valid_frames = tf.reduce_sum(y_mask) * 2.0
    total_valid_frames = tf.maximum(total_valid_frames, 1.0)
    mse = tf.reduce_sum(masked_squared_error) / total_valid_frames
    rmse = tf.sqrt(mse + 1e-8)
    rmse = tf.where(tf.math.is_finite(rmse), rmse, tf.constant(1.0, dtype=rmse.dtype))
    return rmse