Resort alpha-sorted numpy matrix to cluster-sort order

resorty_numpy_2_cluster_sort_order.py

#!/usr/bin/env python


import numpy as np
import time
from datetime import timedelta

def cluster_sort(genus1, genus2):
    start_time = time.monotonic()

    n1 = 'names1.txt'
    n2 = 'names2.txt'
    nn1 = 'new_names1.txt'
    nn2 = 'new_names2.txt'
    distance_matrix = 'distance.npy'
    sorted_matrix = 'sorted.npy'

    orig_mtx = np.load(distance_matrix)



    names1_dict = {}
    names2_dict = {}
    new_names1_dict = {}
    new_names2_dict = {}

    new_mtx = np.zeros(orig_mtx.size, dtype=np.float64)
    new_mtx = new_mtx.ravel()

    with open(n1, 'r') as names1, open(n2, 'r') as names2, open(nn1, 'r') as new_names1, open(nn2, 'r') as new_names2 :
        i = 0
        for line in names1:
            names1_dict[line.strip()] = i
            i = i + 1

        i = 0
        for line in names2:
            names2_dict[line.strip()] = i
            i = i + 1

        i = 0
        for line in new_names1:
            new_names1_dict[i] = line.strip()
            i = i + 1

        i = 0
        for line in new_names2:
            new_names2_dict[i] = line.strip()
            i = i + 1

    rows = orig_mtx.shape[0]
    cols = orig_mtx.shape[1]

    for x in range(rows):
        for y in range(cols):
            val = orig_mtx[names1_dict[new_names1_dict[x]]][names2_dict[new_names2_dict[y]]]
            new_mtx[x * cols + y] = val

    new_mtx = np.reshape(new_mtx, orig_mtx.shape)
    np.save(sorted_matrix, new_mtx)

    # print("\nNew matrix\n---------------------")
    # print(new_mtx[0])
    # print(new_mtx[1])
    # print(new_mtx[2])
    # print(new_mtx[3])
    # print("\n---------------------\n")

    end_time = time.monotonic()
    print("\nSorting Complete! " + str(timedelta(seconds=end_time - start_time)))

def main():
    parser = argparse.ArgumentParser()

    parser.add_argument('-g1', dest='genus1', help='Genus 1', required=True)
    parser.add_argument('-g2', dest='genus2', help='Genus 2', required=True)

    args = parser.parse_args()

    cluster_sort(args.genus1, args.genus2)

    return 0

if __name__ == "__main__":
    sys.exit(main())