"""Analyze condition number of the network."""

import numpy as np
import matplotlib.pyplot as plt

# import model

def _get_sparse_mask(nx, ny, non, complex=False, nOR=50):
    """Generate a binary mask.

    The mask will be of size (nx, ny)
    For all the nx connections to each 1 of the ny units, only non connections are 1.

    Args:
        nx: int
        ny: int
        non: int, must not be larger than nx

    Return:
        mask: numpy array (nx, ny)
    """
    mask = np.zeros((nx, ny))

    if not complex:
        mask[:non] = 1
        for i in range(ny):
            np.random.shuffle(mask[:, i])  # shuffling in-place
    return mask.astype(np.float32)

def _get_cond(q, n_orn, n_pn, n_kc, n_kc_claw):
    M = np.random.rand(n_orn, n_pn)
    M_new = M * (1-q) + np.eye(n_orn) * q
    
    # J = np.random.rand(N_PN, N_KC) / np.sqrt(N_PN + N_KC)
    # J = np.random.randn(N_PN, N_KC) / np.sqrt(N_PN + N_KC)
    J = np.random.rand(n_pn, n_kc)
    mask = _get_sparse_mask(n_pn, n_kc, n_kc_claw) / n_kc_claw
    J = J * mask

    K = np.dot(M_new, J)
    
    # cond = np.linalg.cond(K)

    cond = np.linalg.norm(np.linalg.pinv(K)) * np.linalg.norm(K)
    return cond


def get_logcond(q=1, n_orn=50, n_pn=50, n_kc=2500, n_kc_claw=7, n_rep=10):
    conds = [_get_cond(q, n_orn, n_pn, n_kc, n_kc_claw) for i in range(n_rep)]
    return np.mean(np.log10(conds))



def plot_cond_by_q(n_kc=2500):
    qs = np.linspace(0, 1, 100)
    conds = [get_logcond(q=q, n_kc=n_kc) for q in qs]
    
    plt.figure()
    plt.plot(qs, conds, 'o-')
    plt.title('N_KC: ' + str(n_kc))
    plt.xlabel('fraction diagonal')
    plt.ylabel('log condition number')
    # plt.savefig('figures/condvsfracdiag_nkc'+str(n_kc)+'.pdf', transparent=True)


def plot_cond_by_n_kc():
    n_kcs = np.logspace(1, 4, 10).astype(int)
    conds_q1 = np.array([get_logcond(n_kc=n_kc, q=1) for n_kc in n_kcs])
    
    plt.figure()
    plt.plot(np.log10(n_kcs), conds_q1, 'o-')
    plt.xticks(np.log10(n_kcs), n_kcs)
    plt.xlabel('N_KC')
    
    
    n_kcs = np.logspace(1, 4, 10).astype(int)
    conds_q0 = np.array([get_logcond(n_kc=n_kc, q=0) for n_kc in n_kcs])
    
    plt.figure()
    plt.plot(np.log10(n_kcs), conds_q0, 'o-')
    plt.xticks(np.log10(n_kcs), n_kcs)
    plt.xlabel('N_KC')
    
    plt.figure()
    plt.plot(np.log10(n_kcs), conds_q1 - conds_q0, 'o-')
    plt.xticks(np.log10(n_kcs), n_kcs)
    plt.ylabel('Log decrease in condition number')
    plt.xlabel('N_KC')
    
    
n_kc_claws = np.arange(1, 50)
conds = np.array([get_logcond(n_kc_claw=n) for n in n_kc_claws])

plt.figure()
plt.plot(n_kc_claws, conds, 'o-')
plt.xticks(n_kc_claws)
plt.xlabel('N_KC_claw')
plt.show()

[((2495, 2511), 'numpy.arange', 'np.arange', (['(1)', '(50)'], {}), '(1, 50)\n', (2504, 2511), True, 'import numpy as np\n'), ((2578, 2590), 'matplotlib.pyplot.figure', 'plt.figure', ([], {}), '()\n', (2588, 2590), True, 'import matplotlib.pyplot as plt\n'), ((2591, 2624), 'matplotlib.pyplot.plot', 'plt.plot', (['n_kc_claws', 'conds', '"""o-"""'], {}), "(n_kc_claws, conds, 'o-')\n", (2599, 2624), True, 'import matplotlib.pyplot as plt\n'), ((2625, 2647), 'matplotlib.pyplot.xticks', 'plt.xticks', (['n_kc_claws'], {}), '(n_kc_claws)\n', (2635, 2647), True, 'import matplotlib.pyplot as plt\n'), ((2648, 2671), 'matplotlib.pyplot.xlabel', 'plt.xlabel', (['"""N_KC_claw"""'], {}), "('N_KC_claw')\n", (2658, 2671), True, 'import matplotlib.pyplot as plt\n'), ((2672, 2682), 'matplotlib.pyplot.show', 'plt.show', ([], {}), '()\n', (2680, 2682), True, 'import matplotlib.pyplot as plt\n'), ((485, 503), 'numpy.zeros', 'np.zeros', (['(nx, ny)'], {}), '((nx, ny))\n', (493, 503), True, 'import numpy as np\n'), ((732, 759), 'numpy.random.rand', 'np.random.rand', (['n_orn', 'n_pn'], {}), '(n_orn, n_pn)\n', (746, 759), True, 'import numpy as np\n'), ((936, 962), 'numpy.random.rand', 'np.random.rand', (['n_pn', 'n_kc'], {}), '(n_pn, n_kc)\n', (950, 962), True, 'import numpy as np\n'), ((1052, 1068), 'numpy.dot', 'np.dot', (['M_new', 'J'], {}), '(M_new, J)\n', (1058, 1068), True, 'import numpy as np\n'), ((1422, 1444), 'numpy.linspace', 'np.linspace', (['(0)', '(1)', '(100)'], {}), '(0, 1, 100)\n', (1433, 1444), True, 'import numpy as np\n'), ((1508, 1520), 'matplotlib.pyplot.figure', 'plt.figure', ([], {}), '()\n', (1518, 1520), True, 'import matplotlib.pyplot as plt\n'), ((1525, 1550), 'matplotlib.pyplot.plot', 'plt.plot', (['qs', 'conds', '"""o-"""'], {}), "(qs, conds, 'o-')\n", (1533, 1550), True, 'import matplotlib.pyplot as plt\n'), ((1591, 1622), 'matplotlib.pyplot.xlabel', 'plt.xlabel', (['"""fraction diagonal"""'], {}), "('fraction diagonal')\n", (1601, 1622), True, 'import matplotlib.pyplot as plt\n'), ((1627, 1661), 'matplotlib.pyplot.ylabel', 'plt.ylabel', (['"""log condition number"""'], {}), "('log condition number')\n", (1637, 1661), True, 'import matplotlib.pyplot as plt\n'), ((1900, 1912), 'matplotlib.pyplot.figure', 'plt.figure', ([], {}), '()\n', (1910, 1912), True, 'import matplotlib.pyplot as plt\n'), ((2002, 2020), 'matplotlib.pyplot.xlabel', 'plt.xlabel', (['"""N_KC"""'], {}), "('N_KC')\n", (2012, 2020), True, 'import matplotlib.pyplot as plt\n'), ((2159, 2171), 'matplotlib.pyplot.figure', 'plt.figure', ([], {}), '()\n', (2169, 2171), True, 'import matplotlib.pyplot as plt\n'), ((2261, 2279), 'matplotlib.pyplot.xlabel', 'plt.xlabel', (['"""N_KC"""'], {}), "('N_KC')\n", (2271, 2279), True, 'import matplotlib.pyplot as plt\n'), ((2289, 2301), 'matplotlib.pyplot.figure', 'plt.figure', ([], {}), '()\n', (2299, 2301), True, 'import matplotlib.pyplot as plt\n'), ((2402, 2448), 'matplotlib.pyplot.ylabel', 'plt.ylabel', (['"""Log decrease in condition number"""'], {}), "('Log decrease in condition number')\n", (2412, 2448), True, 'import matplotlib.pyplot as plt\n'), ((2453, 2471), 'matplotlib.pyplot.xlabel', 'plt.xlabel', (['"""N_KC"""'], {}), "('N_KC')\n", (2463, 2471), True, 'import matplotlib.pyplot as plt\n'), ((1153, 1170), 'numpy.linalg.norm', 'np.linalg.norm', (['K'], {}), '(K)\n', (1167, 1170), True, 'import numpy as np\n'), ((1362, 1377), 'numpy.log10', 'np.log10', (['conds'], {}), '(conds)\n', (1370, 1377), True, 'import numpy as np\n'), ((1926, 1941), 'numpy.log10', 'np.log10', (['n_kcs'], {}), '(n_kcs)\n', (1934, 1941), True, 'import numpy as np\n'), ((1974, 1989), 'numpy.log10', 'np.log10', (['n_kcs'], {}), '(n_kcs)\n', (1982, 1989), True, 'import numpy as np\n'), ((2185, 2200), 'numpy.log10', 'np.log10', (['n_kcs'], {}), '(n_kcs)\n', (2193, 2200), True, 'import numpy as np\n'), ((2233, 2248), 'numpy.log10', 'np.log10', (['n_kcs'], {}), '(n_kcs)\n', (2241, 2248), True, 'import numpy as np\n'), ((2315, 2330), 'numpy.log10', 'np.log10', (['n_kcs'], {}), '(n_kcs)\n', (2323, 2330), True, 'import numpy as np\n'), ((2374, 2389), 'numpy.log10', 'np.log10', (['n_kcs'], {}), '(n_kcs)\n', (2382, 2389), True, 'import numpy as np\n'), ((588, 619), 'numpy.random.shuffle', 'np.random.shuffle', (['mask[:, (i)]'], {}), '(mask[:, (i)])\n', (605, 619), True, 'import numpy as np\n'), ((784, 797), 'numpy.eye', 'np.eye', (['n_orn'], {}), '(n_orn)\n', (790, 797), True, 'import numpy as np\n'), ((1132, 1149), 'numpy.linalg.pinv', 'np.linalg.pinv', (['K'], {}), '(K)\n', (1146, 1149), True, 'import numpy as np\n'), ((1784, 1805), 'numpy.logspace', 'np.logspace', (['(1)', '(4)', '(10)'], {}), '(1, 4, 10)\n', (1795, 1805), True, 'import numpy as np\n'), ((2043, 2064), 'numpy.logspace', 'np.logspace', (['(1)', '(4)', '(10)'], {}), '(1, 4, 10)\n', (2054, 2064), True, 'import numpy as np\n')]