Merge remote-tracking branch 'official/hydrolysis_2' into hydrolysis_2

arc/species/converter.py

@@ -51,6 +51,9 @@
 logger = get_logger()
 
 
+global BEST_GUESSES
+
+
 def str_to_xyz(xyz_str: str,
                project_directory: Optional[str] = None,
                ) -> dict:
@@ -2180,7 +2183,7 @@ def add_atom_to_xyz_using_internal_coords(xyz: Union[dict, str],
                                           r_value: float,
                                           a_value: float,
                                           d_value: float,
-                                          opt_method: str = 'SLSQP',
+                                          opt_methods: Optional[Union[str, List[str]]] = None,
                                           ) -> Optional[dict]:
     """
     Add an atom to an XYZ structure based on distance, angle, and dihedral constraints.
@@ -2197,13 +2200,17 @@ def add_atom_to_xyz_using_internal_coords(xyz: Union[dict, str],
         r_value (float): The value of the R-X distance parameter, r.
         a_value (float): The value of the A-B-X angle parameter, a.
         d_value (float): The value of the L-M-N-X dihedral angle parameter, d.
-        opt_method (str, optional): The optimization method to use for finding the new atom's coordinates.
-                                    Additional options include 'Nelder-Mead', 'trust-constr' and 'BFGS'.
+        opt_methods (List[str], optional): The optimization method to use for finding the new atom's coordinates.
+                                           Options include 'SLSQP', 'Nelder-Mead', 'trust-constr' and 'BFGS'.
 
     Returns:
         Optional[dict]: The updated xyz coordinates.
     """
+    global BEST_GUESSES
+    BEST_GUESSES = list()
     xyz = check_xyz_dict(xyz)
+    opt_methods = [opt_methods] if isinstance(opt_methods, str) else opt_methods
+    opt_methods = opt_methods or ['SLSQP', 'Nelder-Mead', 'BFGS']
 
     def calculate_errors(result_coords):
         """Calculate constraint errors for distance, angle, and dihedral"""
@@ -2223,39 +2230,69 @@ def meets_precision(result_coords):
         generate_perpendicular_initial_guess,
         generate_bond_length_initial_guess,
         generate_shifted_initial_guess,
+        average_best_guesses_from_all_methods,
     ]
 
     best_error = float('inf')
+    best_guess_and_errors = dict()
+
+    for opt_method in opt_methods:
+        best_guess = None
+        for guess_func in guess_functions:
+            initial_guess = guess_func(
+                atom_r_coord=xyz['coords'][r_index],
+                r_value=r_value,
+                atom_a_coord=xyz['coords'][a_indices[0]],
+                atom_b_coord=xyz['coords'][a_indices[1]],
+                a_value=a_value,
+            )
 
-    for guess_func in guess_functions:
-        initial_guess = guess_func(
-            atom_r_coord=xyz['coords'][r_index],
-            r_value=r_value,
-            atom_a_coord=xyz['coords'][a_indices[0]],
-            atom_b_coord=xyz['coords'][a_indices[1]],
-            a_value=a_value
-        )
-
-        try:
-            updated_xyz = _add_atom_to_xyz_using_internal_coords(
-                xyz, element, r_index, a_indices, d_indices, r_value, a_value, d_value, initial_guess, opt_method)
-            new_coord = updated_xyz['coords'][-1]
-            r_error, a_error, d_error = calculate_errors(new_coord)
-            total_error = r_error + a_error + d_error
-
-            if meets_precision(new_coord):
-                print("Precision criteria met. Returning result.")
-                return updated_xyz
-
-            if total_error < best_error:
-                print(f"Updating closest result. Previous best_error={best_error}, new total_error={total_error}")
-                best_error = total_error
-
-        except Exception as e:
-            print(f"Optimization failed with {guess_func.__name__}: {e}")
-
-    print("No valid solution was found.")
-    return None
+            try:
+                updated_xyz = _add_atom_to_xyz_using_internal_coords(
+                    xyz, element, r_index, a_indices, d_indices, r_value, a_value, d_value, initial_guess, opt_method)
+                new_coord = updated_xyz['coords'][-1]
+                r_error, a_error, d_error = calculate_errors(new_coord)
+                total_error = r_error + a_error + d_error
+
+                if meets_precision(new_coord):
+                    print("Precision criteria met. Returning result.")
+                    return updated_xyz
+
+                if total_error < best_error:
+                    print(f"Updating closest result. Previous best_error={best_error}, new total_error={total_error}")
+                    best_error = total_error
+                    best_guess = updated_xyz
+                    best_guess_and_errors[best_error] = best_guess
+
+            except Exception as e:
+                print(f"Optimization failed with {guess_func.__name__}: {e}")
+        BEST_GUESSES.append(best_guess)
+
+    print("No valid solution was found, returning the closest solution after brute force modifications.")
+    lowest_key = None
+    for key, val in best_guess_and_errors.items():
+        if lowest_key is None or key < lowest_key:
+            lowest_key = key
+    xyz = best_guess_and_errors[lowest_key]
+    xyz = modify_coords(coords=xyz,
+                        indices=[r_index, len(xyz['coords']) - 1],
+                        new_value=r_value,
+                        modification_type='atom',
+                        index=0,
+                        )
+    xyz = modify_coords(coords=xyz,
+                        indices=[a_indices[0], a_indices[1], len(xyz['coords']) - 1],
+                        new_value=a_value,
+                        modification_type='atom',
+                        index=0,
+                        )
+    xyz = modify_coords(coords=xyz,
+                        indices=[d_indices[0], d_indices[1], d_indices[2], len(xyz['coords']) - 1],
+                        new_value=d_value,
+                        modification_type='atom',
+                        index=0,
+                        )
+    return xyz
 
 
 def _add_atom_to_xyz_using_internal_coords(xyz: dict,
@@ -2513,8 +2550,15 @@ def generate_shifted_initial_guess(atom_r_coord, r_value, atom_a_coord, atom_b_c
     base_guess = generate_initial_guess_r_a(atom_r_coord, r_value, atom_a_coord, atom_b_coord, a_value)
     return base_guess + shift
 
+def average_best_guesses_from_all_methods(atom_r_coord, r_value, atom_a_coord, atom_b_coord, a_value):
+    global BEST_GUESSES
+    if not BEST_GUESSES:
+        return generate_initial_guess_r_a(atom_r_coord, r_value, atom_a_coord, atom_b_coord, a_value)
+    coords = np.array([xyz['coords'][-1] for xyz in BEST_GUESSES])
+    return np.mean(coords, axis=0)
+
 def generate_bond_length_initial_guess(atom_r_coord, r_value, atom_a_coord, atom_b_coord, a_value):
     """Generate an initial guess considering only the bond length to the reference atom."""
     direction = np.random.uniform(-1.0, 1.0, 3)  # Random direction
     direction /= np.linalg.norm(direction)  # Normalize to unit vector
-    return np.array(atom_r_coord) + r_value * direction
+    return np.array(atom_r_coord) + r_value * direction