Skip to content
This repository was archived by the owner on May 5, 2026. It is now read-only.

Commit ee9ceaf

Browse files
ss0832ss0832
authored
Add files via upload
1 parent 1616124 commit ee9ceaf

1 file changed

Lines changed: 314 additions & 0 deletions

File tree

‎multioptpy/OtherMethod/spring_pair_method.py‎

Lines changed: 314 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,314 @@
1+
import numpy as np
2+
import os
3+
import shutil
4+
5+
from multioptpy.Utils.calc_tools import Calculationtools
6+
from multioptpy.Visualization.visualization import Graph
7+
8+
class SpringPairMethod:
9+
"""
10+
Implementation of the Spring Pair Method (SPM) with Adaptive Step Size & Momentum.
11+
Modified to accept a SINGLE initial structure and auto-generate the second one via perturbation.
12+
"""
13+
def __init__(self, config):
14+
self.config = config
15+
16+
# --- TUNING PARAMETERS ---
17+
self.k_spring = 10.0
18+
self.l_s = max(getattr(self.config, 'L_covergence', 0.1), 0.1)
19+
self.initial_drift_step = 0.01
20+
self.climb_step_size = 0.50
21+
self.drift_limit = 100
22+
self.momentum = 0.3
23+
self.MAX_FORCE_THRESHOLD = 0.00100
24+
self.RMS_FORCE_THRESHOLD = 0.00005
25+
26+
# Magnitude of random perturbation to generate Image 2 (Bohr)
27+
self.perturbation_scale = 0.1
28+
29+
def RMS(self, mat):
30+
return np.sqrt(np.mean(mat**2))
31+
32+
def print_info(self, dat, phase):
33+
print(f"[[{phase} information]]")
34+
print(f" image_1 image_2")
35+
print(f"energy (Hartree) : {dat['energy_1']:.8f} {dat['energy_2']:.8f}")
36+
print(f"gradient RMS : {self.RMS(dat['gradient_1']):.6f} {self.RMS(dat['gradient_2']):.6f}")
37+
38+
if "perp_force_1" in dat:
39+
print(f"perp_force (Drift) : {self.RMS(dat['perp_force_1']):.6f} {self.RMS(dat['perp_force_2']):.6f}")
40+
print(f"spring_force : {self.RMS(dat['spring_force_1']):.6f} {self.RMS(dat['spring_force_2']):.6f}")
41+
42+
print(f"spring_length (Bohr) : {dat['spring_length']:.6f}")
43+
step_info = dat.get('step_size', 0)
44+
print(f"DEBUG: k={self.k_spring:.1f}, step={step_info:.6f}")
45+
print("-" * 80)
46+
return
47+
48+
def get_spring_vectors(self, geom_1, geom_2):
49+
diff = geom_2 - geom_1
50+
dist = np.linalg.norm(diff)
51+
if dist < 1e-10:
52+
rand_vec = np.random.randn(*diff.shape)
53+
unit_vec = rand_vec / np.linalg.norm(rand_vec)
54+
dist = 1e-10
55+
else:
56+
unit_vec = diff / dist
57+
return dist, unit_vec
58+
59+
def decompose_gradient(self, gradient, unit_vec):
60+
grad_flat = gradient.flatten()
61+
vec_flat = unit_vec.flatten()
62+
grad_par_mag = np.dot(grad_flat, vec_flat)
63+
grad_par = grad_par_mag * unit_vec
64+
grad_perp = gradient - grad_par
65+
return grad_par, grad_perp
66+
67+
def _generate_perturbed_structure(self, geom, scale):
68+
"""Generate a new structure by adding random perturbation to the given geometry."""
69+
# Generate random noise
70+
noise = np.random.randn(*geom.shape)
71+
# Normalize to unit vectors per atom to distribute perturbation evenly
72+
norms = np.linalg.norm(noise, axis=1, keepdims=True)
73+
noise = noise / (norms + 1e-10)
74+
# Scale the noise
75+
perturbation = noise * scale
76+
return geom + perturbation
77+
78+
def iteration(self, file_directory_1, SP1, element_list, electric_charge_and_multiplicity, FIO1):
79+
"""
80+
Main SPM Optimization Loop.
81+
Accepts ONE initial structure. Image 2 is auto-generated.
82+
"""
83+
G = Graph(self.config.iEIP_FOLDER_DIRECTORY)
84+
ENERGY_LIST_1, ENERGY_LIST_2 = [], []
85+
GRAD_LIST_1, GRAD_LIST_2 = [], []
86+
87+
# --- Initialize Image 2 from Image 1 ---
88+
print("### Initializing SPM ###")
89+
print(f"Base Structure (Image 1): {file_directory_1}")
90+
91+
# 1. Get initial geometry of Image 1
92+
# We perform a dummy single point calculation or just read the geom if possible.
93+
# Here we run SP1 to get the geometry and energy reliably.
94+
init_energy, init_grad, init_geom, err = SP1.single_point(
95+
file_directory_1, element_list, "init_check",
96+
electric_charge_and_multiplicity, self.config.force_data["xtb"]
97+
)
98+
if err:
99+
print("[Error] Failed to read initial structure.")
100+
return
101+
102+
# 2. Generate Image 2 geometry
103+
print(f"Generating Image 2 with random perturbation (Scale: {self.perturbation_scale} Bohr)...")
104+
init_geom_2 = self._generate_perturbed_structure(init_geom, self.perturbation_scale)
105+
106+
# Create a directory for Image 2 (internally managed)
107+
# We can just use the same SP1 object but we need a file path for it.
108+
# We will generate input files for Image 2 on the fly using FIO1 logic.
109+
# Note: We reuse SP1 and FIO1 for both images since they share physics/settings.
110+
111+
# Initialize loop variables
112+
geom_num_list_1 = init_geom
113+
geom_num_list_2 = init_geom_2
114+
115+
velocity_1 = np.zeros((len(element_list), 3))
116+
velocity_2 = np.zeros((len(element_list), 3))
117+
118+
current_drift_step = self.initial_drift_step
119+
120+
# Variables to store the latest geometry
121+
new_geom_1 = None
122+
new_geom_2 = None
123+
124+
for cycle in range(0, self.config.microiterlimit):
125+
if os.path.isfile(self.config.iEIP_FOLDER_DIRECTORY+"end.txt"):
126+
break
127+
128+
print(f"### Cycle {cycle} Start ###")
129+
130+
# =========================================================================
131+
# 1. Drifting Phase
132+
# =========================================================================
133+
print(f"--- Drifting Phase (Cycle {cycle}) ---")
134+
135+
prev_force_drift_1 = None
136+
prev_force_drift_2 = None
137+
138+
drift_temp_label = f"{cycle}_drift_temp"
139+
140+
for d_step in range(self.drift_limit):
141+
iter_label = drift_temp_label
142+
143+
# Make input files for this step
144+
# Note: file_directory_1/2 here are just strings for the input file path
145+
# generated by _make_next_input.
146+
# However, SP1.single_point expects the directory path or file path depending on implementation.
147+
# Assuming _make_next_input returns the PATH to the input file.
148+
149+
input_path_1 = self._make_next_input(FIO1, geom_num_list_1, element_list, electric_charge_and_multiplicity, iter_label + "_img1")
150+
input_path_2 = self._make_next_input(FIO1, geom_num_list_2, element_list, electric_charge_and_multiplicity, iter_label + "_img2")
151+
152+
# 1. QM Calculation
153+
# We reuse SP1 for both. It's just a calculator.
154+
energy_1, gradient_1, g1_read, error_flag_1 = SP1.single_point(input_path_1, element_list, iter_label + "_img1", electric_charge_and_multiplicity, self.config.force_data["xtb"])
155+
energy_2, gradient_2, g2_read, error_flag_2 = SP1.single_point(input_path_2, element_list, iter_label + "_img2", electric_charge_and_multiplicity, self.config.force_data["xtb"])
156+
157+
# Update geom from read result to be safe, or stick to internal numpy array
158+
# Using internal array (geom_num_list_1/2) is safer for consistency unless optimizer changes it.
159+
# But let's align them just in case output orientation changed.
160+
# g1_read, g2_read = Calculationtools().kabsch_algorithm(g1_read, g2_read)
161+
# geom_num_list_1, geom_num_list_2 = g1_read, g2_read
162+
163+
# Align current internal geometries
164+
geom_num_list_1, geom_num_list_2 = Calculationtools().kabsch_algorithm(geom_num_list_1, geom_num_list_2)
165+
166+
if error_flag_1 or error_flag_2:
167+
return
168+
169+
# 2. Vector & Force Calculation
170+
ds, vs = self.get_spring_vectors(geom_num_list_1, geom_num_list_2)
171+
172+
_, grad_perp_1 = self.decompose_gradient(gradient_1, vs)
173+
_, grad_perp_2 = self.decompose_gradient(gradient_2, vs)
174+
175+
force_perp_1 = -grad_perp_1
176+
force_perp_2 = -grad_perp_2
177+
178+
spring_force_mag = self.k_spring * (ds - self.l_s)
179+
spring_force_1 = spring_force_mag * vs
180+
spring_force_2 = spring_force_mag * (-vs)
181+
182+
total_force_1 = force_perp_1 + spring_force_1
183+
total_force_2 = force_perp_2 + spring_force_2
184+
185+
# --- Adaptive Step Logic ---
186+
if prev_force_drift_1 is not None:
187+
dot_1 = np.sum(prev_force_drift_1 * total_force_1)
188+
dot_2 = np.sum(prev_force_drift_2 * total_force_2)
189+
190+
if dot_1 < 0 or dot_2 < 0:
191+
current_drift_step *= 0.5
192+
velocity_1 *= 0.1
193+
velocity_2 *= 0.1
194+
print(f" [Auto-Brake] Oscillation detected at step {d_step}. Reduced drift step to {current_drift_step:.6f}")
195+
else:
196+
current_drift_step = min(current_drift_step * 1.05, self.initial_drift_step)
197+
198+
prev_force_drift_1 = total_force_1.copy()
199+
prev_force_drift_2 = total_force_2.copy()
200+
201+
# Update Position
202+
velocity_1 = self.momentum * velocity_1 + current_drift_step * total_force_1
203+
velocity_2 = self.momentum * velocity_2 + current_drift_step * total_force_2
204+
205+
geom_num_list_1 += velocity_1
206+
geom_num_list_2 += velocity_2
207+
208+
# Check Convergence
209+
drift_metric = max(self.RMS(force_perp_1), self.RMS(force_perp_2))
210+
211+
if d_step % 5 == 0:
212+
info_dat = {
213+
"energy_1": energy_1, "energy_2": energy_2,
214+
"gradient_1": gradient_1, "gradient_2": gradient_2,
215+
"perp_force_1": force_perp_1, "perp_force_2": force_perp_2,
216+
"spring_force_1": spring_force_1, "spring_force_2": spring_force_2,
217+
"spring_length": ds,
218+
"step_size": current_drift_step,
219+
"convergence_metric": drift_metric
220+
}
221+
self.print_info(info_dat, f"Cycle {cycle} - Drifting {d_step}")
222+
223+
if drift_metric < self.RMS_FORCE_THRESHOLD:
224+
print(f" >> Drifting converged at step {d_step}")
225+
break
226+
227+
# =========================================================================
228+
# 2. Climbing Phase
229+
# =========================================================================
230+
print(f"--- Climbing Phase (Cycle {cycle}) ---")
231+
232+
iter_label = f"{cycle}_climb"
233+
234+
input_path_1 = self._make_next_input(FIO1, geom_num_list_1, element_list, electric_charge_and_multiplicity, iter_label + "_img1")
235+
input_path_2 = self._make_next_input(FIO1, geom_num_list_2, element_list, electric_charge_and_multiplicity, iter_label + "_img2")
236+
237+
energy_1, gradient_1, g1_read, error_flag_1 = SP1.single_point(input_path_1, element_list, iter_label + "_img1", electric_charge_and_multiplicity, self.config.force_data["xtb"])
238+
energy_2, gradient_2, g2_read, error_flag_2 = SP1.single_point(input_path_2, element_list, iter_label + "_img2", electric_charge_and_multiplicity, self.config.force_data["xtb"])
239+
240+
geom_num_list_1, geom_num_list_2 = Calculationtools().kabsch_algorithm(geom_num_list_1, geom_num_list_2)
241+
242+
ds, vs = self.get_spring_vectors(geom_num_list_1, geom_num_list_2)
243+
grad_par_1, _ = self.decompose_gradient(gradient_1, vs)
244+
grad_par_2, _ = self.decompose_gradient(gradient_2, vs)
245+
246+
active_climb_step = self.climb_step_size
247+
248+
# Move UP
249+
geom_num_list_1 += active_climb_step * grad_par_1
250+
geom_num_list_2 += active_climb_step * grad_par_2
251+
252+
# Update 'new_geom' for final output reference
253+
new_geom_1 = geom_num_list_1.copy()
254+
new_geom_2 = geom_num_list_2.copy()
255+
256+
grad_norm_1 = np.linalg.norm(gradient_1)
257+
grad_norm_2 = np.linalg.norm(gradient_2)
258+
global_metric = min(grad_norm_1, grad_norm_2)
259+
260+
info_dat = {
261+
"energy_1": energy_1, "energy_2": energy_2,
262+
"gradient_1": gradient_1, "gradient_2": gradient_2,
263+
"par_force_1": grad_par_1, "par_force_2": grad_par_2,
264+
"spring_length": ds, "convergence_metric": global_metric,
265+
"step_size": active_climb_step
266+
}
267+
self.print_info(info_dat, f"Cycle {cycle} - Climbing")
268+
269+
ENERGY_LIST_1.append(energy_1 * self.config.hartree2kcalmol)
270+
ENERGY_LIST_2.append(energy_2 * self.config.hartree2kcalmol)
271+
GRAD_LIST_1.append(grad_norm_1)
272+
GRAD_LIST_2.append(grad_norm_2)
273+
274+
if cycle > 5 and global_metric < self.MAX_FORCE_THRESHOLD:
275+
print("!!! Global Convergence Reached !!!")
276+
print(f"Saddle point candidate found around cycle {cycle}")
277+
break
278+
279+
# --- END OF OPTIMIZATION LOOP ---
280+
281+
# Save the optimized structure
282+
if new_geom_1 is not None and new_geom_2 is not None:
283+
avg_geom = (new_geom_1 + new_geom_2) / 2.0
284+
avg_geom_ang = avg_geom * self.config.bohr2angstroms
285+
286+
dir_name = os.path.basename(os.path.normpath(self.config.iEIP_FOLDER_DIRECTORY))
287+
output_xyz_name = f"{dir_name}_optimized.xyz"
288+
289+
try:
290+
with open(output_xyz_name, "w") as f:
291+
num_atoms = len(element_list)
292+
f.write(f"{num_atoms}\n")
293+
f.write(f"SPM Optimized Saddle Point (Average)\n")
294+
for i, elem in enumerate(element_list):
295+
x, y, z = avg_geom_ang[i]
296+
f.write(f"{elem} {x:.10f} {y:.10f} {z:.10f}\n")
297+
print(f"\n[Success] Final optimized structure saved to: {output_xyz_name}")
298+
except Exception as e:
299+
print(f"\n[Error] Failed to save optimized structure: {e}")
300+
301+
return
302+
303+
def _make_next_input(self, FIO, geom, element_list, charge_mult, iter_label):
304+
"""
305+
Creates a PSI4 input file and returns its path.
306+
"""
307+
geom_tolist = (geom * self.config.bohr2angstroms).tolist()
308+
for i, elem in enumerate(element_list):
309+
geom_tolist[i].insert(0, elem)
310+
geom_tolist.insert(0, charge_mult)
311+
312+
# FIO.make_psi4_input_file usually returns the directory or path.
313+
# Ensure this matches your FIO implementation.
314+
return FIO.make_psi4_input_file([geom_tolist], iter_label)

0 commit comments

Comments
 (0)