Skip to content

Commit

Permalink
temp
Browse files Browse the repository at this point in the history
  • Loading branch information
@JintaoWu98
JintaoWu98 committed Dec 5, 2024
1 parent f14c6e8 commit 1e96bd6
Show file tree
Hide file tree
Showing 3 changed files with 208 additions and 184 deletions.
141 changes: 68 additions & 73 deletions arc/plotter.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1016,11 +1016,14 @@ def plot_torsion_angles(torsion_angles,
num_comb = None
torsions = list(torsion_angles.keys()) if torsions_sampling_points is None \
else list(torsions_sampling_points.keys())
print("torsions_sampling_points:", torsions_sampling_points)
print("torsion_angles:", torsion_angles)
ticks = [0, 60, 120, 180, 240, 300, 360]
sampling_points = dict()
if torsions_sampling_points is not None:
for tor, points in torsions_sampling_points.items():
sampling_points[tor] = [point if point <= 360 else point - 360 for point in points]
if not isinstance(points[0],list):
sampling_points[tor] = [point if point <= 360 else point - 360 for point in points]
if not torsions:
return
if len(torsions) == 1:
Expand Down Expand Up @@ -1049,6 +1052,11 @@ def plot_torsion_angles(torsion_angles,
fig.dpi = 120
num_comb = 1
for i, torsion in enumerate(torsions):
torsion_key = tuple(torsion)
print(f"Checking torsion key: {torsion_key}")
if torsion_key not in torsion_angles:
print(f"Key {torsion_key} not found in torsion_angles. Skipping...")
continue
axs[i].plot(np.array(torsion_angles[tuple(torsion)]),
np.zeros_like(np.arange(len(torsion_angles[tuple(torsion)]))), 'g.')
if wells_dict is not None:
Expand Down Expand Up @@ -1122,88 +1130,75 @@ def plot_torsion_angles(torsion_angles,
return num_comb


def plot_ring_torsion_angles(conformers, plot_path=None):
def plot_ring_torsion_angles(conformers, plot_path=None, tolerance=15):
"""
Plot the torsion angles of the generated conformers.
Plot the torsion angles of the generated conformers for each ring,
considering torsion similarity within a given tolerance.
Args:
torsion_angles (dict): Keys are torsions, values are lists of corresponding angles.
torsions_sampling_points (dict, optional): Keys are torsions, values are sampling points.
wells_dict (dict, optional): Keys are torsions, values are lists of wells.
Each entry in such a list is a well dictionary with the following keys:
``start_idx``, ``end_idx``, ``start_angle``, ``end_angle``, and ``angles``.
e_conformers (list, optional): Entries are conformers corresponding to the sampling points with FF energies.
de_threshold (float, optional): Energy threshold, plotted as a dashed horizontal line.
conformers (list): A list of conformers, each containing a 'puckering' key with ring torsion angles.
plot_path (str, optional): The path for saving the plot.
tolerance (float, optional): The angular tolerance to consider two torsion angle sets as similar.
"""
num_comb = None
if 'puckering' not in conformers[0]:
return
ring_num = len(conformers[0]['puckering'])
# fig, axs = plt.subplots(nrows=ring_num, ncols=1, sharex=True, sharey=True, gridspec_kw={'hspace': 0})
# fig.dpi = 120
num_comb = 1
for i in range(ring_num):
unique_angle_sets = {}
for conformer in conformers:
angles = conformer['puckering']
for ring, angle_set in angles.items():
# Round angles to the nearest integer
rounded_angle_set = tuple(round(angle) for angle in angle_set)
found = False
for existing_angle_set in unique_angle_sets.keys():
if almost_equal_lists(existing_angle_set, rounded_angle_set):
unique_angle_sets[existing_angle_set] += 1
found = True
break
if not found:
unique_angle_sets[rounded_angle_set] = 1
angles = [list(angles) for angles in unique_angle_sets.keys()]
counts = list(unique_angle_sets.values())

# Combine angles and counts into a single list of tuples for sorting
angles_counts = list(zip(angles, counts))

# Sort by counts in descending order
angles_counts_sorted = sorted(angles_counts, key=lambda x: x[1], reverse=True)

# Unzip the sorted tuples back into separate lists
angles_sorted, counts_sorted = zip(*angles_counts_sorted)

# Dictionary to store unique angle sets for each ring
ring_angle_data = {}

fig, ax = plt.subplots(figsize=(10, 5))
x = np.arange(len(angles_sorted)) # the label locations
# Process each conformer
for conformer in conformers:
rings = conformer['puckering'] # Retrieve the puckering angles for rings
for torsions, angle_set in rings.items():
rounded_angle_set = tuple(round(angle) for angle in angle_set) # Round angles
if torsions not in ring_angle_data:
ring_angle_data[torsions] = []

# Check for similarity within the current ring
is_similar = False
for i, (existing_set, count) in enumerate(ring_angle_data[torsions]):
if all(abs(a1 - a2) <= tolerance for a1, a2 in zip(rounded_angle_set, existing_set)):
# If similar, increment count
ring_angle_data[torsions][i] = (existing_set, count + 1)
is_similar = True
break
if not is_similar:
# Add unique angle set with a count
ring_angle_data[torsions].append((rounded_angle_set, 1))


ax.bar(x, counts_sorted, color='blue')
ax.set_xlabel('Base Angles Set (Rounded to Nearest Integer)')
ax.set_ylabel('Frequency')
ax.set_title('Frequency of Different Angle Sets Across Conformers')
ax.set_xticks(x)
ax.set_xticklabels([f'{angle}' for angle in angles_sorted], rotation=45, ha='right')
# Plot data for each ring
for ring, angle_counts in ring_angle_data.items():
# Extract and sort data
angles, counts = zip(*angle_counts)
angles_counts_sorted = sorted(zip(angles, counts), key=lambda x: x[1], reverse=True)
angles_sorted, counts_sorted = zip(*angles_counts_sorted)

# Create bar plot for this ring
fig, ax = plt.subplots(figsize=(10, 5))
x = np.arange(len(angles_sorted)) # Label positions
ax.bar(x, counts_sorted, color='blue')
ax.set_xlabel('Rounded Angle Sets (Nearest Integer)')
ax.set_ylabel('Frequency')
ax.set_title(f'Frequency of Different Angle Sets for Ring {ring}')
ax.set_xticks(x)
ax.set_xticklabels([f'{angle}' for angle in angles_sorted], rotation=45, ha='right')

# Save or display the plot
if plot_path is not None:
ring_plot_path = os.path.join(plot_path, f'conformer_ring_torsions_{ring}.png')
if not os.path.isdir(plot_path):
os.makedirs(plot_path)
try:
plt.savefig(ring_plot_path, bbox_inches='tight')
except FileNotFoundError:
pass
if is_notebook():
plt.show()
plt.close(fig)

# plt.tight_layout()
# plt.show()
# plt.setp(axs, xticks=ticks) # set the x ticks of all subplots
# fig.set_size_inches(8, len(torsions) * 1.5)
if plot_path is not None:
if not os.path.isdir(plot_path):
os.makedirs(plot_path)
file_names = list()
for (_, _, files) in os.walk(plot_path):
file_names.extend(files)
break # don't continue to explore subdirectories
i = 0
for file_ in file_names:
if 'conformer torsions' in file_:
i += 1
image_path = os.path.join(plot_path, f'conformer ring torsions {i}.png')
try:
plt.savefig(image_path, bbox_inches='tight')
except FileNotFoundError:
pass
if is_notebook():
plt.show()
plt.close(fig)
return num_comb
# Return the collected data
return ring_angle_data


def plot_1d_rotor_scan(angles: Optional[Union[list, tuple, np.array]] = None,
Expand Down
53 changes: 41 additions & 12 deletions arc/species/conformers.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -98,7 +98,7 @@
MAX_COMBINATION_ITERATIONS = 25

# A threshold below which all combinations will be generated. Above it just samples of the entire search space.
COMBINATION_THRESHOLD = 1000
COMBINATION_THRESHOLD = 10

# Consolidation tolerances for Z matrices
CONSOLIDATION_TOLS = {'R': 1e-2, 'A': 1e-2, 'D': 1e-2}
Expand Down Expand Up @@ -352,19 +352,29 @@ def deduce_new_conformers(label, conformers, torsions, tops, mol_list, smeared_s
symmetries[tuple(torsion)] = symmetry
logger.debug(f'Identified {len([s for s in symmetries.values() if s > 1])} symmetric wells for {label}')

torsions_sampling_points, wells_dict = dict(), dict()
torsions_sampling_points, wells_dict, ring_sampling_points = dict(), dict(), dict()
for tor, tor_angles in torsion_angles.items():
torsions_sampling_points[tor], wells_dict[tor] = \
determine_torsion_sampling_points(label, tor_angles, smeared_scan_res=smeared_scan_res,
symmetry=symmetries[tor])

print('wells_dict:', wells_dict)
if plot_path is not None:
arc.plotter.plot_torsion_angles(torsion_angles, torsions_sampling_points, wells_dict=wells_dict,
plot_path=plot_path)
arc.plotter.plot_ring_torsion_angles(conformers=conformers, plot_path=plot_path)
angles_sorted = arc.plotter.plot_ring_torsion_angles(conformers=conformers, plot_path=plot_path)
# Process the ring angle data
for ring, angle_counts in angles_sorted.items():
# Extract unique angles only
angles = [list(angle) for angle, _ in angle_counts] # Ignore counts, only keep the angles

# Flatten the structure and add to torsions_sampling_points
ring_sampling_points[tuple(ring)] = angles
print('ring_sampling_points:', ring_sampling_points)

combined_sampling_points = {**torsions_sampling_points, **ring_sampling_points}

hypothetical_num_comb = 1
for points in torsions_sampling_points.values():
for points in combined_sampling_points.values():
hypothetical_num_comb *= len(points)
number_of_chiral_centers = get_number_of_chiral_centers(label, mol, conformer=conformers[0],
just_get_the_number=True)
Expand All @@ -375,10 +385,10 @@ def deduce_new_conformers(label, conformers, torsions, tops, mol_list, smeared_s
hypothetical_num_comb_str = str(hypothetical_num_comb)
logger.info(f'\nHypothetical number of conformer combinations for {label}: {hypothetical_num_comb_str}')

# split torsions_sampling_points into two lists, use combinations only for those with multiple sampling points
# split combined_sampling_points into two lists, use combinations only for those with multiple sampling points
single_tors, multiple_tors, single_sampling_point, multiple_sampling_points = list(), list(), list(), list()
multiple_sampling_points_dict = dict() # used for plotting an energy "scan"
for tor, points in torsions_sampling_points.items():
for tor, points in combined_sampling_points.items():
if len(points) == 1:
single_tors.append(tor)
single_sampling_point.append((points[0]))
Expand Down Expand Up @@ -541,7 +551,8 @@ def conformers_combinations_by_lowest_conformer(label, mol, base_xyz, multiple_t
'FF energy': round(energy, 3),
'source': f'Changing dihedrals on most stable conformer, iteration {i}',
'torsion': tor,
'dihedral': round(dihedral, 2),
'dihedral': round(dihedral, 2) if isinstance(dihedral, float)
else [round(angle, 2) for angle in dihedral],
'dmat': dmat1,
'fl_distance': fl_distance1}
newest_conformers_dict[tor].append(conformer)
Expand All @@ -557,7 +568,7 @@ def conformers_combinations_by_lowest_conformer(label, mol, base_xyz, multiple_t
'FF energy': None,
'source': f'Changing dihedrals on most stable conformer, iteration {i}, but FF energy is None',
'torsion': tor,
'dihedral': round(dihedral, 2),
'dihedral': round(dihedral, 2) if isinstance(dihedral, float) else [round(angle, 2) for angle in dihedral],
'dmat': dmat1,
'fl_distance': fl_distance1})
new_conformers.extend(newest_conformer_list)
Expand All @@ -572,6 +583,7 @@ def conformers_combinations_by_lowest_conformer(label, mol, base_xyz, multiple_t
if plot_path is not None:
logger.info(converter.xyz_to_str(lowest_conf_i['xyz']))
arc.plotter.draw_structure(xyz=lowest_conf_i['xyz'])
print('hhhh')
num_comb = arc.plotter.plot_torsion_angles(torsion_angles, multiple_sampling_points_dict,
wells_dict=wells_dict, e_conformers=newest_conformers_dict,
de_threshold=de_threshold, plot_path=plot_path)
Expand Down Expand Up @@ -760,8 +772,26 @@ def change_dihedrals_and_force_field_it(label, mol, xyz, torsions, new_dihedrals
for torsion, dihedral in zip(torsions, dihedrals):
conf, rd_mol = converter.rdkit_conf_from_mol(mol, xyz_dihedrals)
if conf is not None:
torsion_0_indexed = [tor - 1 for tor in torsion]
xyz_dihedrals = converter.set_rdkit_dihedrals(conf, rd_mol, torsion_0_indexed, deg_abs=dihedral)
if isinstance(torsion[0], int): # Assume torsion is a flat tuple
torsion_0_indexed = [tor - 1 for tor in torsion]
elif isinstance(torsion[0], tuple): # Check if torsion is a tuple of tuples
torsion_0_indexed = [[tor - 0 for tor in sub_torsion] for sub_torsion in torsion]
print("change_torsion", torsion)
print("change_dihedral", dihedral)
if not isinstance(dihedral, list):
xyz_dihedrals = converter.set_rdkit_dihedrals(conf, rd_mol, torsion_0_indexed, deg_abs=dihedral)
elif isinstance(dihedral, list):
# Calculate the head and tail
ring_length = len(torsion_0_indexed) # Length of the ring
head = torsion_0_indexed[0][0] # First element of the first torsion
# Find the tail
for torsion in torsion_0_indexed:
if head == torsion[-1]: # Check if the head appears in the torsion
print("torsion", torsion)
tail = torsion[-2] # The last element in this torsion becomes the tail candidate
break
print(f'we are here ring twist, head {head}, tail {tail}')
xyz_dihedrals = converter.set_rdkit_ring_dihedrals(conf, rd_mol, head, tail, torsion_0_indexed[0:ring_length-3], dihedral[0:ring_length-3])
xyz_, energy = get_force_field_energies(label, mol=mol, xyz=xyz_dihedrals, optimize=True,
force_field=force_field, suppress_warning=True)
if energy and xyz_:
Expand Down Expand Up @@ -926,7 +956,6 @@ def determine_puckering(conformers, rings_indexes):
list: Entries are conformer dictionaries.
"""
for conformer in conformers:
# print(f'conformer:{conformer}')
if isinstance(conformer['xyz'], str):
xyz = converter.str_to_xyz(conformer['xyz'])
else:
Expand Down
Loading

0 comments on commit 1e96bd6

Please sign in to comment.