print 'Calculating rates for coupling ' + str(V)
    rates = []
    for i, delta_E in enumerate(delta_E_values):
        site_hamiltonian = hamiltonian(delta_E, V) + np.diag(
            total_site_reorg_energies
        )  # adjust site energies by reorganisation shift
        evals, evecs = utils.sorted_eig(site_hamiltonian)
        # calculate exciton reorganisation energies
        exciton_reorg_energies = np.array([
            os.exciton_reorg_energy(evecs[i], total_site_reorg_energies)
            for i in range(site_hamiltonian.shape[0])
        ])
        print exciton_reorg_energies
        evals = evals - exciton_reorg_energies  # adjust to bare exciton reorganisation energies
        MRT = os.modified_redfield_rates(evals, evecs, g_site, g_site_dot,
                                         g_site_dot_dot,
                                         total_site_reorg_energy, temperature,
                                         time)
        rates.append(MRT[0, 1])
    rates_data.append(rates)

#np.savez('../../data/modified_redfield_test_high_energy_modes_data.npz', delta_E_values=delta_E_values, coupling_values=coupling_values, rates=rates_data)
for i, rates in enumerate(rates_data):
    plt.subplot(1, 3, i + 1)
    plt.plot(delta_E_values, utils.WAVENUMS_TO_INVERSE_PS * np.array(rates))
plt.show()

# data = np.load('../../data/modified_redfield_test_high_energy_modes_data.npz')
# rates = data['rates']
# delta_E_values = data['delta_E_values']
# plt.plot(delta_E_values, -utils.WAVENUMS_TO_INVERSE_PS* rates)
# plt.show()
num_expansion_terms = 20
g_site, g_site_dot, g_site_dot_dot, total_site_reorg_energy = os.modified_redfield_params(time, site_reorg_energy, cutoff_freq, temperature, mode_params, num_expansion_terms)
total_site_reorg_energies = np.array([total_site_reorg_energy, total_site_reorg_energy])
rates_data = []
   
print 'Calculating rates with high energy modes....'
for V in coupling_values:
    print 'Calculating rates for coupling ' + str(V)
    rates = []
    for i,delta_E in enumerate(delta_E_values):
        site_hamiltonian = hamiltonian(delta_E, V) + np.diag(total_site_reorg_energies) # adjust site energies by reorganisation shift
        evals, evecs = utils.sorted_eig(site_hamiltonian)
        # calculate exciton reorganisation energies
        exciton_reorg_energies = np.array([os.exciton_reorg_energy(evecs[i], total_site_reorg_energies) for i in range(site_hamiltonian.shape[0])])
        print exciton_reorg_energies
        evals = evals - exciton_reorg_energies # adjust to bare exciton reorganisation energies
        MRT = os.modified_redfield_rates(evals, evecs, g_site, g_site_dot, g_site_dot_dot, total_site_reorg_energy, temperature, time)
        rates.append(MRT[0,1])
    rates_data.append(rates)
      
#np.savez('../../data/modified_redfield_test_high_energy_modes_data.npz', delta_E_values=delta_E_values, coupling_values=coupling_values, rates=rates_data)
for i,rates in enumerate(rates_data):
    plt.subplot(1,3,i+1)
    plt.plot(delta_E_values, utils.WAVENUMS_TO_INVERSE_PS*np.array(rates))
plt.show()

# data = np.load('../../data/modified_redfield_test_high_energy_modes_data.npz')
# rates = data['rates']
# delta_E_values = data['delta_E_values']
# plt.plot(delta_E_values, -utils.WAVENUMS_TO_INVERSE_PS* rates)
# plt.show()
 '''
 # remove reorg energy and reorder evals/evecs from lowest to highest energy
 evals, evecs = utils.sort_evals_evecs(evals-exciton_reorg_energies, evecs)
 
 # calculate modified Redfield rates
 MRT_rates = os.modified_redfield_rates_general(evals, evecs, lbf, lbf_dot, lbf_dot_dot, total_site_reorg_energy, temperature, time)[0]
 
 
 # need to reorder exciton_reorg_energies too! easier to just recalculate with new evecs
 exciton_reorg_energies = np.array([os.exciton_reorg_energy(exciton, site_reorg_energies) for exciton in evecs])
 '''
 
 '''
 temporary to test that this reproduces the results for incorrect ordering of exciton levels
 '''
 MRT_rates = os.modified_redfield_rates(evals-exciton_reorg_energies, evecs, lbf, lbf_dot, lbf_dot_dot, total_site_reorg_energy, temperature, time)
 
 
 site_lbfs = np.array([lbf, lbf, lbf, lbf, lbf, lbf])
 exciton_lbfs = np.array([os.exciton_lbf(exciton, site_lbfs) for exciton in evecs])
  
 primary_CT_scaling = 3.
 primary_CT_reorg_energy = primary_CT_scaling * total_site_reorg_energy
 primary_CT_state = np.array([0, 0, 0, 0, 0, 0, 1., 0])
 secondary_CT_scaling = 4.
 secondary_CT_reorg_energy = secondary_CT_scaling * total_site_reorg_energy
 secondary_CT_state = np.array([0, 0, 0, 0, 0, 0, 0, 1.])
  
 # calculate Forster rates to primary and secondary CT states
 forward_primary_forster_rates = np.zeros(site_hamiltonian.shape[0])
 backward_primary_forster_rates = np.zeros(site_hamiltonian.shape[0])