#!/usr/bin/env python3

import matplotlib.pyplot as plt
import scipy.linalg as la
import numpy as np

def Main():

    V = 0

    x = [2**n for n in range(3, 12)]
    y = []

    print('N\tS')
    for N in x:

        s2, S = Solution(N, V)

        y.append(S)

        print(f'{N}\t{S:.2f}')

    p = np.polyfit(np.log(x), y, 1)
    plt.plot(x, [p[0]*i + p[1] for i in np.log(x)])

    plt.semilogx(x, y, '.')
    plt.show()

    return

def Solution(N, V):
    H = np.zeros((N+1, N+1))

    ### fill in the Hamiltonian here ###
    ### example:
    for i in range(N):
        H[i, i] = 1
    H[-1, -1] = 1
    ### H == I

    _, V = la.eigh(H, subset_by_index=[0, 0])

    ### Ground State
    psi = V[:, 0]

    ### Calculate singular values (squared) here ###
    s2 = ...

    S = ...

    return s2, S

Main()