fix issues regarding random search

cagrikymk · Feb 16, 2024 · 2813449 · 2813449
1 parent 0956396
commit 2813449
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/jaxreaxff/driver.py b/jaxreaxff/driver.py
@@ -339,7 +339,7 @@ def new_loss_func(params, param_indices,
     print("Trial-{} is starting...".format(i+1))
     start = time.time()
     if args.init_FF_type == 'random':
-      [min_params, min_loss] = random_parameter_search(bounds, random_sample_count,
+      min_params = random_parameter_search(bounds, random_sample_count,
                                   param_indices, force_field, training_data,
                                   list_positions, aligned_data, center_sizes,
                                   new_loss_func)

diff --git a/jaxreaxff/optimizer.py b/jaxreaxff/optimizer.py
@@ -539,7 +539,7 @@ def random_parameter_search(bounds, sample_count,
   for _ in range(sample_count):
     selected_params = onp.random.uniform(low=bounds[:,0],high=bounds[:,1])
     selected_params = jnp.array(selected_params, dtype=dtype)
-    loss = loss_func(selected_params, args)
+    loss = loss_func(selected_params, *args)
     if loss < min_loss or onp.isnan(min_loss) == True:
       min_loss = loss
       min_params = selected_params