overlap_v6.py

#!/usr/bin/env python3

"""This script reads the eddy graph from Matlab v6 files
    id_child_(anti|cyclo).mat and N(anti|cyclo).mat, and writes the
    graph to a new file in edgelist format.

"""

import scipy.io as sio
import csv
import sys

orientation = sys.argv[1] # "cyclo" or "anti"
id_child = sio.loadmat(f"id_child_{orientation}.mat", squeeze_me=True)\
    ["id_child"]
n_eddies = sio.loadmat(f"N{orientation}.mat", squeeze_me = True)\
    [f"N{orientation}"].astype(int, casting = "safe", copy = False)
n_dates = n_eddies.size
print("n_dates =", n_dates)
e_overestim = id_child.shape[0]
print("e_overestim =", e_overestim)

with open(f"edgelist_{orientation}.csv", "w", newline = '') as edgelist:
    writer = csv.writer(edgelist, delimiter = ' ', lineterminator = "\n")

    for k1 in range(n_dates):
        for i1, id_child_1 in enumerate(id_child[:n_eddies[k1], k1]):
            n1 = k1 * e_overestim + i1 + 1
            # id_child_1 may be an int, which means there is a single
            # successor, or id_child_1 may be a numpy array, which means
            # there are several successors.

            try:
                # Try processing id_child_1 as an array:
                for n2 in id_child_1:
                    writer.writerow((n1, n2))
            except TypeError:
                # id_child_1 is a scalar int.
                writer.writerow((n1, id_child_1))