print('\nRotation Matrix') print(np.array(diabatization['rot_matrix'])) print('\nAdiabatic Matrix') print(np.array(diabatization['adiabatic_matrix'])) print('\nDiabatic Matrix') print(np.array(diabatization['diabatic_matrix'])) print('\nDiabatic states dimer\n--------------------') from pyqchem.tools import plot_diabatization plot_diabatization(diabatization['diabatic_states'], atoms_ranges=[ dimer.get_number_of_atoms() / 2, dimer.get_number_of_atoms() ]) # Monomer adiabatic states (extra test) qc_input = QchemInput( opt_monomer, jobtype='sp', exchange='hf', correlation='rasci', basis='sto-3g', ras_act=6, ras_elec=4, ras_spin_mult=1, # singlets only ras_roots=4, # calculate 8 states ras_do_hole=False,
def parser_rasci(output): """ Parser for RAS-CI calculations Include: - Diabatization scheme data - Structure - Adiabatic states - SOC :param output: :return: """ data_dict = {} # Molecule n = output.find('$molecule') n2 = output[n:].find('$end') molecule_region = output[n:n + n2 - 1].replace('\t', ' ').split('\n')[1:] charge, multiplicity = [int(num) for num in molecule_region[0].split()] coordinates = [[float(l) for l in line.split()[1:4]] for line in molecule_region[1:]] symbols = [line.split()[0].capitalize() for line in molecule_region[1:]] n_atoms = len(symbols) # structure structure_input = Structure(coordinates=coordinates, symbols=symbols, charge=charge, multiplicity=multiplicity) enum = output.find('Standard Nuclear Orientation') section_structure = output[enum:enum + 200 * structure_input.get_number_of_atoms()].split( '\n') section_structure = section_structure[3:structure_input. get_number_of_atoms() + 3] coordinates = [[float(num) for num in s.split()[2:]] for s in section_structure] data_dict['structure'] = Structure(coordinates=coordinates, symbols=symbols, charge=charge, multiplicity=multiplicity) # basic info enum = output.find('Nuclear Repulsion Energy') basic_data = read_basic_info(output[enum:enum + 5000]) # scf_energy enum = output.find('SCF energy in the final basis set') scf_energy = float(output[enum:enum + 100].split()[8]) data_dict['scf_energy'] = scf_energy # total energy # enum = output.find('Total energy in the final basis set') # total_energy = float(output[enum:enum+100].split()[8]) # RASCI dimensions ini_section = output.find('RAS-CI Dimensions') end_section = search_bars(output, from_position=enum, bar_type='\*\*\*')[1] dimension_section = output[ini_section:end_section] enum = dimension_section.find('Doubly Occ') doubly_occ = int(dimension_section[enum:enum + 50].split()[3]) enum = dimension_section.find('Doubly Vir') doubly_vir = int(dimension_section[enum:enum + 50].split()[2]) enum = dimension_section.find('Frozen Occ') frozen_occ = int(dimension_section[enum:enum + 50].split()[3]) enum = dimension_section.find('Frozen Vir') frozen_vir = int(dimension_section[enum:enum + 50].split()[2]) enum = dimension_section.find('Total CI configurations') total_conf = int(dimension_section[enum:enum + 50].split()[3]) enum = dimension_section.find('Active configurations') active_conf = int(dimension_section[enum:enum + 50].split()[2]) enum = dimension_section.find('Hole configurations') hole_conf = int(dimension_section[enum:enum + 50].split()[2]) enum = dimension_section.find('Particle configurations') particle_conf = int(dimension_section[enum:enum + 50].split()[2]) rasci_dimensions = { 'doubly_occupied': doubly_occ, 'doubly_virtual': doubly_vir, 'frozen_occupied': frozen_occ, 'frozen_virtual': frozen_vir, 'total_configurations': total_conf, 'active_configurations': active_conf, 'hole_configurations': hole_conf, 'particle_configurations': particle_conf } data_dict.update({'rasci_dimensions': rasci_dimensions}) # Diabatization scheme done_diabat = bool(output.find('RASCI DIABATIZATION') + 1) if done_diabat: rot_matrix = _read_simple_matrix( 'showmatrix final adiabatic -> diabatic', output)[-1] adiabatic_matrix = _read_simple_matrix( 'showing H in adiabatic representation: NO coupling elements', output)[-1] diabatic_matrix = _read_simple_matrix( 'showing H in diabatic representation: WITH coupling elements', output)[-1] mulliken_adiabatic = [] enum = output.find('Mulliken analysis of Adiabatic State') for m in re.finditer('Mulliken analysis of Adiabatic State', output[enum:]): section_mulliken = output[m.end() + enum:m.end() + 10000 + enum] # 10000: assumed to max of section section_mulliken = section_mulliken[:section_mulliken.find( 'Natural Orbitals stored in FCHK')] section_attachment = section_mulliken.split('\n')[9 + n_atoms:9 + n_atoms * 2] mulliken_adiabatic.append({ 'attach': [float(l.split()[1]) for l in section_attachment], 'detach': [float(l.split()[2]) for l in section_attachment], 'total': [float(l.split()[3]) for l in section_attachment] }) mulliken_diabatic = [] enum = output.find('showing H in diabatic representation') for m in re.finditer('Mulliken Analysis of Diabatic State', output[enum:]): section_mulliken = output[m.end() + enum:m.end() + 10000 + enum] # 10000: assumed to max of section section_mulliken = section_mulliken[:section_mulliken.find( 'Natural Orbitals stored in FCHK')] section_attachment = section_mulliken.split('\n')[9 + n_atoms:9 + n_atoms * 2] mulliken_diabatic.append({ 'attach': [float(l.split()[1]) for l in section_attachment], 'detach': [float(l.split()[2]) for l in section_attachment], 'total': [float(l.split()[3]) for l in section_attachment] }) enum = output.find('Transition dipole moment - diabatic states') tdm_section = output[enum:enum + 70 * len(rot_matrix)] diabatic_tdm = [] for m in re.finditer('TDM', tdm_section): diabatic_tdm.append([ float(n) for n in tdm_section[m.end():m.end() + 70][14:].split()[:3] ]) diabatic_states = [] for i, tdm in enumerate(diabatic_tdm): diabatic_states.append({ 'excitation_energy': diabatic_matrix[i][i], 'excitation_energy_units': 'eV', 'transition_moment': tdm, 'dipole_moment_units': 'ua', 'mulliken': mulliken_diabatic[i] }) data_dict['diabatization'] = { 'rot_matrix': rot_matrix, 'adiabatic_matrix': adiabatic_matrix, 'diabatic_matrix': diabatic_matrix, 'diabatic_states': diabatic_states, 'mulliken_adiabatic': mulliken_adiabatic } # excited states data excited_states = [] for m in re.finditer('RAS-CI total energy for state', output): # print('ll found', m.start(), m.end()) section_state = output[m.end():m.end() + 10000] # 10000: assumed to max of section section_state = section_state[:section_state.find('********')] enum = section_state.find('RAS-CI total energy for state') section_state = section_state[:enum] # energies tot_energy = float(section_state.split()[1]) exc_energy_units = section_state.split()[4][1:-1] exc_energy = float(section_state.split()[6]) state_multiplicity = section_state.split( )[8] if section_state.split()[8] != ':' else section_state.split()[9] # dipole moment enum = section_state.find('Dipole Moment') dipole_mom = [ float(section_state[enum:].split()[2]) + 0.0, float(section_state[enum:].split()[4]) + 0.0, float(section_state[enum:].split()[6]) + 0.0 ] # Transition moment enum = section_state.find('Trans. Moment') if enum > -1: trans_mom = [ float(section_state[enum:].split()[2]) + 0.0, float(section_state[enum:].split()[4]) + 0.0, float(section_state[enum:].split()[6]) + 0.0 ] trans_mom = standardize_vector(trans_mom) strength = float(section_state[enum:].split()[10]) else: trans_mom = None strength = None # configurations table enum = section_state.find('AMPLITUDE') enum2 = section_state.find('Contributions') section_table = section_state[enum:enum2].split('\n')[2:-2] # ' HOLE | ALPHA | BETA | PART | AMPLITUDE' table = [] for row in section_table: table.append({ 'hole': row.split('|')[1].strip(), 'alpha': row.split('|')[2].strip(), 'beta': row.split('|')[3].strip(), 'part': row.split('|')[4].strip(), 'amplitude': float(row.split('|')[5]) + 0.0 }) table[-1]['occupations'] = get_rasci_occupations_list( table[-1], data_dict['structure'], basic_data['n_basis_functions']) table = sorted(table, key=operator.itemgetter('hole', 'alpha', 'beta', 'part')) # Contributions RASCI wfn contributions_section = section_state[enum2:] contributions = { 'active': float(contributions_section.split()[4]), 'hole': float(contributions_section.split()[6]), 'part': float(contributions_section.split()[8]) } # complete dictionary tot_energy_units = 'au' excited_states.append({ 'total_energy': tot_energy, 'total_energy_units': tot_energy_units, 'excitation_energy': exc_energy, 'excitation_energy_units': exc_energy_units, 'multiplicity': state_multiplicity, 'dipole_moment': dipole_mom, 'transition_moment': trans_mom, 'dipole_moment_units': 'ua', 'oscillator_strength': strength, 'configurations': table, 'contributions_fwn': contributions }) data_dict.update({'excited_states': excited_states}) # Interstate transition properties done_interstate = bool(output.find('Interstate Transition Properties') + 1) if done_interstate: ini_section = output.find('Interstate Transition Properties') end_section = search_bars(output, from_position=ini_section)[1] interstate_section = output[ini_section:end_section] interstate_dict = {} for m in re.finditer('State A: Root', interstate_section): section_pair = interstate_section[m.start():m.start() + 10000] section_pair = section_pair[:section_pair.find('********')] lines = section_pair.split('\n') state_a = int(lines[0].split()[-1]) state_b = int(lines[1].split()[-1]) pair_dict = {'state_a': state_a, 'state_b': state_b} s_a = s_b = 0 for i, line in enumerate(lines): # RAS-CI SOC section if 'Angular momentum components' in line: pair_dict['angular_momentum'] = [ complex(lines[i + 1 + k].split()[-1].replace('i', 'j')) for k in range(3) ] if '||gamma^AB||_total' in line: pair_dict['gamma_total'] = float(lines[i + 0].split()[-1]) pair_dict['gamma_sym'] = float(lines[i + 1].split()[-1]) pair_dict['gamma_anti_sym'] = float(lines[i + 2].split()[-1]) if "KET: S',Sz'" in line: s_a = float(lines[i].split('=')[1].split()[0]) s_b = float(lines[i + 1].split('=')[1].split()[0]) na = int(2 * s_a + 1) nb = int(2 * s_b + 1) if 'Spin Matrices' in line: spinmat_x = _read_soc_matrix(lines[i + 2:], [nb, na]) spinmat_y = _read_soc_matrix(lines[i + 4 + nb:], [nb, na]) spinmat_z = _read_soc_matrix(lines[i + 6 + 2 * nb:], [nb, na]) pair_dict['spin_matrices'] = [ spinmat_x, spinmat_y, spinmat_z ] if 'Spin matrices Sx, Sy and Sz for states' in line: pair_dict['spin_matrices'] = [np.zeros( (nb, na)).tolist()] * 3 if '1-elec SOC matrix (cm-1)' in line: pair_dict['1e_soc_mat'] = _read_soc_matrix( lines[i + 1:], [nb, na]) pair_dict['1e_socc'] = float(lines[i + 2 + nb].split()[-2:][0]) if '2e-SOMF Reduced matrix elements (cm-1)' in line: r, c = lines[i + 1].split()[-2:] pair_dict['hso_l-'] = float(r) + float(c) * 1j r, c = lines[i + 2].split()[-2:] pair_dict['hso_l0'] = float(r) + float(c) * 1j r, c = lines[i + 3].split()[-2:] pair_dict['hso_l+'] = float(r) + float(c) * 1j if '2-elec mean-field SOC matrix (cm-1)' in line: pair_dict['2e_soc_mat'] = _read_soc_matrix( lines[i + 1:], [nb, na]) if 'Total mean-field SOC matrix (cm-1)' in line: pair_dict['total_soc_mat'] = _read_soc_matrix( lines[i + 1:], [nb, na]) if 'Mean-Field SOCC' in line: pair_dict['mf_socc'] = float(line.split()[-2]) pair_dict['units'] = line.split()[-1] interstate_dict[(state_a, state_b)] = pair_dict data_dict.update({'interstate_properties': interstate_dict}) return data_dict
def basic_cis(output): """ Parser for CIS/TD-DFT calculations :param output: :return: """ data_dict = {} # Molecule n = output.find('$molecule') n2 = output[n:].find('$end') molecule_region = output[n:n + n2 - 1].replace('\t', ' ').split('\n')[1:] charge, multiplicity = [int(num) for num in molecule_region[0].split()] coordinates = [[float(l) for l in line.split()[1:4]] for line in molecule_region[1:]] symbols = [line.split()[0].capitalize() for line in molecule_region[1:]] n_atoms = len(symbols) # structure structure_input = Structure(coordinates=coordinates, symbols=symbols, charge=charge, multiplicity=multiplicity) enum = output.find('Standard Nuclear Orientation') section_structure = output[enum:enum + 200 * structure_input.get_number_of_atoms()].split( '\n') section_structure = section_structure[3:structure_input. get_number_of_atoms() + 3] coordinates = [[float(num) for num in s.split()[2:]] for s in section_structure] data_dict['structure'] = Structure(coordinates=coordinates, symbols=symbols, charge=charge, multiplicity=multiplicity) # scf_energy enum = output.find('Total energy in the final basis set') try: data_dict['scf_energy'] = float(output[enum:enum + 100].split()[8]) except IndexError: pass enum = output.find('Molecular Point Group') if enum > 0: symmetry_data = read_symmetry_info(output[enum:enum + 1000]) enum = output.find('Nuclear Repulsion Energy') basic_data = read_basic_info(output[enum:enum + 5000]) # CIS excited states # enum = output.find('CIS Excitation Energies') try: enum = list(re.finditer('CIS Excitation Energies', output))[-1].end() except IndexError: enum = list(re.finditer('TDDFT/TDA Excitation Energies', output))[-1].end() excited_states = [] if enum > 0: bars = search_bars(output, from_position=enum) output_cis = output[bars[0]:bars[1]] for m in re.finditer('Excited state ', output_cis): state_cis_section = output_cis[m.end():] state_cis_lines = state_cis_section.split('\n') exc_energy = float(state_cis_lines[0].split()[5]) exc_energy_units = state_cis_lines[0].split()[3][1:-1] tot_energy = float(state_cis_lines[1].split()[5]) try: tot_energy_units = state_cis_lines[1].split()[6] mul = state_cis_lines[2].split()[-1] trans_mom = [ float(mom) for mom in [ state_cis_lines[3].split()[2], state_cis_lines[3]. split()[4], state_cis_lines[3].split()[6] ] ] strength = float(state_cis_lines[4].split()[2]) except ValueError: # old version of qchem (< 5.01) state_cis_words = output_cis[m.end():].split() tot_energy_units = 'au' mul = state_cis_words[13] trans_mom = [ float(mom) for mom in [ state_cis_words[16], state_cis_words[18], state_cis_words[20] ] ] strength = float(state_cis_words[24]) transitions = [] for line in state_cis_lines[5:]: if line.find('-->') > 0: origin = int( line.split('>')[0].split('(')[1].split(')')[0]) target = int( line.split('>')[1].split('(')[1].split(')')[0]) amplitude = float(line.split('=')[1].split()[0]) alpha_transitions = [] beta_transitions = [] try: spin = line[21:].split()[3] if spin == 'alpha': alpha_transitions.append({ 'origin': origin, 'target': target + basic_data['n_alpha'] }) elif spin == 'beta': beta_transitions.append({ 'origin': origin, 'target': target + basic_data['n_beta'] }) else: raise ParserError('basic_cis', 'Error reading configurations') transitions.append({ 'origin': origin, 'target': target, 'amplitude': amplitude, 'occupations': get_cis_occupations_list( basic_data['n_basis_functions'], basic_data['n_alpha'], basic_data['n_beta'], alpha_transitions=alpha_transitions, beta_transitions=beta_transitions) }) except (IndexError, ParserError): # This supposes single electron transition alpha_transitions.append({ 'origin': origin, 'target': target + basic_data['n_alpha'] }) transitions.append({ 'origin': origin, 'target': target, 'amplitude': amplitude / np.sqrt(2), 'occupations': get_cis_occupations_list( basic_data['n_basis_functions'], basic_data['n_alpha'], basic_data['n_beta'], alpha_transitions=alpha_transitions, beta_transitions=beta_transitions) }) transitions.append({ 'origin': origin, 'target': target, 'amplitude': amplitude / np.sqrt(2) if mul == 'Singlet' else -amplitude / np.sqrt(2), 'occupations': get_cis_occupations_list( basic_data['n_basis_functions'], basic_data['n_alpha'], basic_data['n_beta'], alpha_transitions=beta_transitions, beta_transitions=alpha_transitions) }) if len(line) < 5: break excited_states.append({ 'total_energy': tot_energy, 'total_energy_units': tot_energy_units, 'excitation_energy': exc_energy, 'excitation_energy_units': exc_energy_units, 'multiplicity': mul, 'transition_moment': standardize_vector(trans_mom), 'strength': strength, 'configurations': transitions }) data_dict['excited_states'] = excited_states # Spin-Orbit coupling initial = output.find( '*********SPIN-ORBIT COUPLING JOB BEGINS HERE*********') final = output.find('*********SOC CODE ENDS HERE*********') data_interstate = {} if initial > 0: soc_section = output[initial:final] def label_states(excited_states): labels = [] ns = 1 nt = 1 for state in excited_states: if state['multiplicity'].lower() == 'singlet': labels.append('S{}'.format(ns)) ns += 1 elif state['multiplicity'].lower() == 'triplet': labels.append('T{}'.format(nt)) nt += 1 else: try: m = float(state['multiplicity']) if abs(m - 1) < 0.1: labels.append('S{}'.format(ns)) ns += 1 if abs(m - 3) < 0.1: labels.append('T{}'.format(nt)) nt += 1 state['multiplicity'] = m except ValueError: raise ParserError('basic_cis', 'State multiplicity error') return labels, nt - 1, ns - 1 labels, n_triplet, n_singlet = label_states(excited_states) for i, label in enumerate(labels): data_interstate[(i + 1, 0)] = { '1e_soc_mat': [0j, 0j, 0j], 'soc_units': 'cm-1' } data_interstate[(0, i + 1)] = { '1e_soc_mat': [0j, 0j, 0j], 'soc_units': 'cm-1' } for j, label2 in enumerate(labels): if (label[0] == 'S' or label2[0] == 'S') and (label[0] != label2[0]): data_interstate[(i + 1, j + 1)] = { '1e_soc_mat': [[0j, 0j, 0j]], 'soc_units': 'cm-1' } elif label[0] == 'T' and label2[0] == 'T': data_interstate[(i + 1, j + 1)] = { '1e_soc_mat': [[0j, 0j, 0j], [0j, 0j, 0j], [0j, 0j, 0j]], 'soc_units': 'cm-1' } elif label[0] == 'S' and label2[0] == 'S': data_interstate[(i + 1, j + 1)] = { '1e_soc_mat': [[0j]], 'soc_units': 'cm-1' } else: raise ParserError('basic_cis', 'State multiplicity error') for i, label in enumerate(labels): for k2, ms2 in enumerate([-1, 0, 1]): for j, label2 in enumerate(labels): if label[0] == 'T': for k, ms in enumerate([-1, 0, 1]): enum = soc_section.find( 'SOC between the {} (ms={}) state and excited triplet states (ms={})' .format(label, ms2, ms)) for line in soc_section[enum:enum + 80 * (n_triplet + 1)].split( '\n'): if len(line.split()) == 0: break if line.split()[0] == '{}(ms={})'.format( label2, ms): data_interstate[( i + 1, j + 1 )]['1e_soc_mat'][k2][k] = _list_to_complex( line.split()[1:4]) data_interstate[( i + 1, j + 1 )]['1e_soc_mat'][k][k2] = _list_to_complex( line.split()[1:4]) data_interstate[( j + 1, i + 1 )]['1e_soc_mat'][k2][k] = _list_to_complex( line.split()[1:4]) data_interstate[( j + 1, i + 1 )]['1e_soc_mat'][k][k2] = _list_to_complex( line.split()[1:4]) break elif label[0] == 'S': for k, ms in enumerate([-1, 0, 1]): enum = soc_section.find( 'SOC between the {} state and excited triplet states (ms={})' .format(label, ms)) for line in soc_section[enum:enum + 80 * (n_triplet + 1)].split( '\n'): if len(line.split()) == 0: break if line.split()[0] == '{}(ms={})'.format( label2, ms): data_interstate[( i + 1, j + 1 )]['1e_soc_mat'][0][k] = _list_to_complex( line.split()[1:4]) data_interstate[( j + 1, i + 1 )]['1e_soc_mat'][0][k] = _list_to_complex( line.split()[1:4]) break else: raise ParserError('basic_cis', 'SOC reading error') enum = soc_section.find( 'SOC between the singlet ground state and excited triplet states (ms={})' .format(ms2)) for line in soc_section[enum:enum + 80 * (n_triplet + 1)].split('\n'): if len(line.split()) == 0: break if line.split()[0] == '{}(ms={})'.format(label, ms2): data_interstate[( i + 1, 0)]['1e_soc_mat'][k2] = _list_to_complex( line.split()[1:4]) data_interstate[( 0, i + 1)]['1e_soc_mat'][k2] = _list_to_complex( line.split()[1:4]) break data_dict['interstate_properties'] = data_interstate # diabatization initial = output.find('Localization Code for CIS excited states') if initial > 0: bars = search_bars(output, from_position=initial) diabat_section = output[initial:bars[0]] def read_diabatization_matrix(label): matrix = [] for m in re.finditer(label, diabat_section): line = diabat_section[m.end():m.end() + 50].split('\n')[0] matrix.append(float(line.split('=')[1])) diabat_dim = int(np.sqrt(len(matrix))) return np.array(matrix).reshape(diabat_dim, diabat_dim).T r_matrix = read_diabatization_matrix('showmatrix adiabatic R-Matrix') rot_matrix = read_diabatization_matrix( 'showmatrix final adiabatic -> diabatic RotMatrix') diabatic_matrix = read_diabatization_matrix( 'showmatrix diabatH') * AU_TO_EV if len(rot_matrix) == 0: # Boys diabatization rot_matrix = read_diabatization_matrix( 'showmatrix Boys adiabatic->diabatic RotMatrix') diabatic_matrix = read_diabatization_matrix( 'showmatrix Boys diabatH') * AU_TO_EV adiabatic_matrix = read_diabatization_matrix( 'showmatrix adiabatH') * AU_TO_EV diabat_data = { 'rot_matrix': rot_matrix, 'adiabatic_matrix': adiabatic_matrix.tolist(), 'diabatic_matrix': diabatic_matrix.tolist() } if diabat_section.find('showmatrix Total_Decomposed_H_diabatic'): tot_decomp_matrix = read_diabatization_matrix( 'showmatrix Total_Decomposed_H_diabatic') * AU_TO_EV decomp_one_matrix = read_diabatization_matrix( 'showmatrix Decomposed_One_diabatic') * AU_TO_EV decomp_j_matrix = read_diabatization_matrix( 'showmatrix Decomposed_J_diabatic') * AU_TO_EV decomp_k_matrix = read_diabatization_matrix( 'showmatrix Decomposed_K_diabatic') * AU_TO_EV diabat_data.update({ 'tot_decomp_matrix': tot_decomp_matrix, 'decomp_one_matrix': decomp_one_matrix.tolist(), 'decomp_j_matrix': decomp_j_matrix.tolist(), 'decomp_k_matrix': decomp_k_matrix.tolist() }) mulliken_diabatic = [] enum = output.find('Mulliken & Loewdin analysis of') for m in re.finditer('Mulliken analysis of TDA State', output[enum:]): section_mulliken = output[m.end() + enum:m.end() + 10000 + enum] # 10000: assumed to max of section section_mulliken = section_mulliken[:section_mulliken.find( 'Natural Orbitals stored in FCHK')] section_attachment = section_mulliken.split('\n')[10 + n_atoms:10 + n_atoms * 2] mulliken_diabatic.append({ 'attach': [float(l.split()[1]) for l in section_attachment], 'detach': [float(l.split()[2]) for l in section_attachment], 'total': [float(l.split()[3]) for l in section_attachment] }) diabatic_states = [] for i in range(len(rot_matrix)): diabat_states_data = { 'excitation_energy': diabatic_matrix[i][i], 'excitation_energy_units': 'eV', 'transition_moment': [], 'dipole_moment_units': 'ua' } if len(mulliken_diabatic) > 0: diabat_states_data['mulliken'] = mulliken_diabatic[i] diabatic_states.append(diabat_states_data) diabat_data['diabatic_states'] = diabatic_states data_dict['diabatization'] = diabat_data return data_dict