danchern97 / tda4atd Goto Github PK

queue = Queue()
number_of_splits = 2
for i, filename in enumerate(tqdm(adj_filenames, desc='Calculating barcodes')):
    barcodes = defaultdict(list)
    adj_matricies = np.load(filename, allow_pickle=True) # samples X 
    print(f"Matricies loaded from: {filename}")
    ntokens = ntokens_array[i*batch_size*DUMP_SIZE : (i+1)*batch_size*DUMP_SIZE]
    splitted = split_matricies_and_lengths(adj_matricies, ntokens, number_of_splits)
    for matricies, ntokens in tqdm(splitted, leave=False):
        p = Process(
            target=subprocess_wrap,
            args=(
                queue,
                get_only_barcodes,
                (matricies, ntokens, dim, lower_bound)
            )
        )
        p.start()
        barcodes_part = queue.get() # block until putted and get barcodes from the queue
        p.join() # release resources
        p.close() # releasing resources of ripser

nworkers=10
pool = multiprocessing.Pool(nworkers)

args = [(matrices, ntokens, dim, lower_bound) for matrices, ntokens in split_matrices]
all_barcodes = pool.starmap(get_only_barcodes, args)

for barcodes_part in all_barcodes:
    barcodes = unite_barcodes(barcodes, barcodes_part)