From 1b77e897dd6e2c7de012cd93b016758e691cc61e Mon Sep 17 00:00:00 2001 From: ajfriedman22 Date: Tue, 19 Nov 2024 14:34:13 -0700 Subject: [PATCH] Fix atom ordering error for test case --- ensemble_md/utils/coordinate_swap.py | 37 +++++++++++++++++----------- 1 file changed, 22 insertions(+), 15 deletions(-) diff --git a/ensemble_md/utils/coordinate_swap.py b/ensemble_md/utils/coordinate_swap.py index a7cfa6c..b913acf 100644 --- a/ensemble_md/utils/coordinate_swap.py +++ b/ensemble_md/utils/coordinate_swap.py @@ -987,26 +987,33 @@ def write_new_file(df_atom_swap, swap, r_swap, line_start, orig_file, new_file, write_line(new_file, orig_file[i], line, atom_num_B, vel, orig_coords[atom_num_A], resnum, new_res_name) # noqa: E501 res_interest_atom += 1 elif (f'{current_element}{current_extra}{current_num}' in atom_order) or (f'{current_element}V{current_extra}{current_num}' in atom_order) or (f'D{current_element}{current_extra}{current_num}' in atom_order): # Atom is in the molecule, but needs swapped AND there are other atoms before it # noqa: E501 - if (f'{current_element}{current_extra}{current_num}' in atom_order): - atom_pos = atom_order.index(f'{current_element}{current_extra}{current_num}') - elif (f'{current_element}V{current_extra}{current_num}' in atom_order): - atom_pos = atom_order.index(f'{current_element}V{current_extra}{current_num}') - elif (f'D{current_element}{current_extra}{current_num}' in atom_order): - atom_pos = atom_order.index(f'D{current_element}{current_extra}{current_num}') - df_select = _get_subset_df(df_interest, atom_order[res_interest_atom]) - if len(df_select.index) == 0: + if line[1] in df_interest['Name'].values: atom_num_B -= 1 atom_num_A += 1 continue else: - for x in range(res_interest_atom, atom_pos): - df_select = _get_subset_df(df_interest, atom_order[x]) - skip_line = _add_or_swap(df_select, new_file, resnum, new_res_name, vel, atom_num_B, orig_coords, skip_line, atom_order[x]) # noqa: E501 - atom_num_B += 1 + if (f'{current_element}{current_extra}{current_num}' in atom_order): + atom_pos = atom_order.index(f'{current_element}{current_extra}{current_num}') + elif (f'{current_element}V{current_extra}{current_num}' in atom_order): + atom_pos = atom_order.index(f'{current_element}V{current_extra}{current_num}') + elif (f'D{current_element}{current_extra}{current_num}' in atom_order): + atom_pos = atom_order.index(f'D{current_element}{current_extra}{current_num}') + df_select = _get_subset_df(df_interest, atom_order[res_interest_atom]) + if len(df_select.index) == 0: + atom_num_B -= 1 + atom_num_A += 1 + continue + else: + for x in range(res_interest_atom, atom_pos): + print(df_interest) + print(atom_order[x]) + df_select = _get_subset_df(df_interest, atom_order[x]) + skip_line = _add_or_swap(df_select, new_file, resnum, new_res_name, vel, atom_num_B, orig_coords, skip_line, atom_order[x]) # noqa: E501 + atom_num_B += 1 + res_interest_atom += 1 + df_select = df_interest[df_interest['Name'] == line[1]] + skip_line = _add_or_swap(df_select, new_file, resnum, new_res_name, vel, atom_num_B, orig_coords, skip_line, atom_order[res_interest_atom]) # noqa: E501 res_interest_atom += 1 - df_select = df_interest[df_interest['Name'] == line[1]] - skip_line = _add_or_swap(df_select, new_file, resnum, new_res_name, vel, atom_num_B, orig_coords, skip_line, atom_order[res_interest_atom]) # noqa: E501 - res_interest_atom += 1 else: print(f'Warning {line} not written') elif resname != old_res_name and prev_resname == old_res_name: # Add dummy atoms at the end of the residue