survival.py

#!/usr/bin/env python3

import scipy.io as sio
import json
import operator
import csv

for orient1, orient2, orient3 in zip(["cyclo", "anti"], ["Cyclo", "Anti"],
                                     ["C", "Antic"]):
    m = sio.loadmat(f"Trajectories_{orient2}_filt_v6.mat", squeeze_me=True)
    traj = {}
    belong = []

    for t in m[f"{orient3}yclonic_Trajectories"]:
        assert t.ndim == 1
        assert isinstance(t[0], int)
        traj[t[0]] = t[1].tolist()

        for id_det in traj[t[0]]:
            belong.append((id_det, t[0]))

    belong.sort(key = operator.itemgetter(0))

    with open(f"traj_{orient1}.json", "w") as f:
        json.dump(traj, f, indent = 0)

    with open(f"belong_{orient1}.csv", "w", newline='') as f:
        writer = csv.writer(f, lineterminator = "\n")
        writer.writerow(("ID_detection", "ID_traj"))
        writer.writerows(belong)