survival.py

#!/usr/bin/env python3

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

for orient in ["cyclo", "anti"]:
    m = sio.loadmat(f"traj_{orient}.mat", squeeze_me=True)
    traj = {}
    belong = []

    for ID_traj, ID_det_array in zip(m["ID_traj"], m["ID_detection"]):
        assert ID_det_array.ndim == 1
        assert isinstance(ID_traj, int)
        traj[ID_traj] = ID_det_array.tolist()

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

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

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

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