aew_to_tcg_statistics.py

import xarray as xr
import pandas as pd
import numpy as np

from sklearn.linear_model import LogisticRegression
from sklearn.metrics import confusion_matrix, accuracy_score, precision_score, recall_score, f1_score

#•••••••••••••••••••••••••••••••••••••••••••••••••
# SETTINGS
#•••••••••••••••••••••••••••••••••••••••••••••••••
START_YEAR = 1950
END_YEAR = 2015
dist_threshold = 250  # km
MODEL = 'VHR'
RESOLUTION = '0.25'

#•••••••••••••••••••••••••••••••••••••••••••••••••
# HAVERSINE DISTANCE
#•••••••••••••••••••••••••••••••••••••••••••••••••
def haversine(lon1, lat1, lon2, lat2):
    R = 6371  

    lon1, lat1, lon2, lat2 = map(np.radians, [lon1, lat1, lon2, lat2])

    dlon = lon2 - lon1
    dlat = lat2 - lat1

    a = np.sin(dlat/2)**2 + np.cos(lat1)*np.cos(lat2)*np.sin(dlon/2)**2
    c = 2*np.arcsin(np.sqrt(a))

    return R * c


#•••••••••••••••••••••••••••••••••••••••••••••••••
# LOAD TC DATA (FIXED)
#•••••••••••••••••••••••••••••••••••••••••••••••••
tc_file = f"/home/kellie.thrower/senior-project/ipsl_data/ICO-{MODEL}_UZ.csv"
df_tc = pd.read_csv(tc_file)

df_tc.columns = df_tc.columns.str.strip()

df_tc = df_tc.dropna(subset=['year','month','day','hour','lon','lat','track_id'])
df_tc['track_id'] = df_tc['track_id'].astype(str)

# normalize longitude
df_tc['lon'] = (df_tc['lon'] + 180) % 360 - 180

# IMPORTANT: create time FIRST
df_tc['time'] = pd.to_datetime(dict(
    year=df_tc['year'].astype(int),
    month=df_tc['month'].astype(int),
    day=df_tc['day'].astype(int),
    hour=df_tc['hour'].astype(int)
))

# filter years
df_tc = df_tc[(df_tc['year'] >= START_YEAR) & (df_tc['year'] <= END_YEAR)]

# TC genesis (true earliest point per storm)
tc_genesis = df_tc.loc[df_tc.groupby('track_id')['time'].idxmin()].reset_index(drop=True)


#•••••••••••••••••••••••••••••••••••••••••••••••••
# STORAGE
#•••••••••••••••••••••••••••••••••••••••••••••••••
all_records = []


#•••••••••••••••••••••••••••••••••••••••••••••••••
# MAIN LOOP OVER YEARS
#•••••••••••••••••••••••••••••••••••••••••••••••••
for year in range(START_YEAR, END_YEAR + 1):

    print(f"Processing {year}...")

    try:
        aew_file = f"/data/kellie.thrower/IPSL_{MODEL}_QTRACK/{year}_JunOct_{RESOLUTION}_deg/AEW_tracks_post_processed_{year}_junOct.nc"
        ds = xr.open_dataset(aew_file)
    except Exception as e:
        print(f"Skipping {year}: {e}")
        continue

    systems = ds.system.values

    # filter TC for year
    tc_year = tc_genesis[tc_genesis['year'] == year]

    tc_lons = tc_year['lon'].values
    tc_lats = tc_year['lat'].values

    #•••••••••••••••••••••••••••••••••••••••••••••••
    # LOOP AEW SYSTEMS
    #•••••••••••••••••••••••••••••••••••••••••••••••
    for system in systems:

        lon = ds.AEW_lon_smooth.sel(system=system).to_pandas()
        lat = ds.AEW_lat_smooth.sel(system=system).to_pandas()

        valid = lon.notna() & lat.notna()
        lon = lon[valid]
        lat = lat[valid]

        if len(lon) < 2:
            continue

        lon = (lon + 180) % 360 - 180

        start_lon, start_lat = lon.iloc[0], lat.iloc[0]
        end_lon, end_lat = lon.iloc[-1], lat.iloc[-1]

        #••••••••••••••••••••••••••••••••••••••••••••
        # TC MATCH
        #••••••••••••••••••••••••••••••••••••••••••••
        if len(tc_lons) == 0:
            min_dist = np.nan
            is_tc = 0
        else:
            dist = haversine(end_lon, end_lat, tc_lons, tc_lats)
            min_dist = np.min(dist)
            is_tc = int(min_dist <= dist_threshold)

        #••••••••••••••••••••••••••••••••••••••••••••
        # STORE
        #••••••••••••••••••••••••••••••••••••••••••••
        all_records.append({
            "year": year,
            "system": system,
            "start_lon": start_lon,
            "start_lat": start_lat,
            "end_lon": end_lon,
            "end_lat": end_lat,
            "min_tc_dist_km": min_dist,
            "tc_match": is_tc
        })


#•••••••••••••••••••••••••••••••••••••••••••••••••
# BUILD ML DATASET
#•••••••••••••••••••••••••••••••••••••••••••••••••
df = pd.DataFrame(all_records)

features = ["start_lon", "start_lat", "end_lon", "end_lat", "min_tc_dist_km"]

df = df.dropna(subset=features + ["tc_match"])

X = df[features]
y = df["tc_match"]


#•••••••••••••••••••••••••••••••••••••••••••••••••
# LOGISTIC REGRESSION
#•••••••••••••••••••••••••••••••••••••••••••••••••
lr = LogisticRegression(max_iter=2000)
lr.fit(X, y)

df["lr_prob"] = lr.predict_proba(X)[:, 1]
df["lr_pred"] = (df["lr_prob"] >= 0.5).astype(int)


#•••••••••••••••••••••••••••••••••••••••••••••••••
# STATISTICS TABLE
#•••••••••••••••••••••••••••••••••••••••••••••••••
tn, fp, fn, tp = confusion_matrix(y, df["lr_pred"]).ravel()

stats = pd.DataFrame([{
    "TP": tp,
    "TN": tn,
    "FP": fp,
    "FN": fn,
    "Accuracy": accuracy_score(y, df["lr_pred"]),
    "Precision": precision_score(y, df["lr_pred"], zero_division=0),
    "Recall": recall_score(y, df["lr_pred"]),
    "F1": f1_score(y, df["lr_pred"])
}])


#•••••••••••••••••••••••••••••••••••••••••••••••••
# OUTPUT
#•••••••••••••••••••••••••••••••••••••••••••••••••

print("\n====================")
print("FINAL RESULTS")
print("====================")
print(stats)

print("\nTotal samples:", len(df))
print("TC match rate:", df["tc_match"].mean()*100,'%')