Repository: OpenSourceEconomics/estimagic

Issue:
Error in multistart with `ls` optimizers
Hi! this is Mateo from last year's DSE Summer School. I'm starting to use estimagic and it is just, well, magical! Thank you!

I've run into this small bug.

### Bug description

Least squares algorithms appear not to work with multistarts. I think the issue is that some updating step of multistarts is looking for the best criterion *value* in a list where the least-squares algo has stored the vector of residuals instead.

### To Reproduce

```
import estimagic as em
import numpy as np

def loss(params):

    resid = np.array([params['x'], params['y']])
    out = {
        "root_contributions": resid,
        "contributions": np.square(resid),
        "value": np.sum(np.square(resid)),
    }
    return out

est = em.minimize(
    criterion=loss,
    params={'x': -1.0, 'y': 1.0},
    lower_bounds={'x': -10, 'y': -10},
    upper_bounds={'x': 10, 'y': 10},
    algorithm='scipy_ls_trf',
    multistart=True,
    multistart_options={
        "n_samples": 3,
        "share_optimizations": 1.0
    }
)

print(est)
```

### Expected behavior

I'd expect the above code to perform three starts.

### Screenshots/Error messages

Instead, it raises the following error after the first start

```
Traceback (most recent call last):
  File "/home/mvg/em_fail.py", line 16, in <module>
    est = em.minimize(
  File "/usr/local/Anaconda3/lib/python3.9/site-packages/estimagic/optimization/optimize.py", line 403, in minimize
    return _optimize(
  File "/usr/local/Anaconda3/lib/python3.9/site-packages/estimagic/optimization/optimize.py", line 766, in _optimize
    raw_res = run_multistart_optimization(
  File "/usr/local/Anaconda3/lib/python3.9/site-packages/estimagic/optimization/tiktak.py", line 168, in run_multistart_optimization
    state, is_converged = update_convergence_state(
  File "/usr/local/Anaconda3/lib/python3.9/site-packages/estimagic/optimization/tiktak.py", line 471, in update_convergence_state
    if valid_new_y[best_index] <= best_y:
ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()
```

### System

 - OS: Ubuntu 18.04


Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/OpenSourceEconomics/estimagic/pull/411?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=OpenSourceEconomics) Report
> Merging [#411](https://codecov.io/gh/OpenSourceEconomics/estimagic/pull/411?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=OpenSourceEconomics) (a416a0a) into [main](https://codecov.io/gh/OpenSourceEconomics/estimagic/commit/737f4d73386d8782f7f2436ea515096177d67a9b?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=OpenSourceEconomics) (737f4d7) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

```diff
@@           Coverage Diff           @@
##             main     #411   +/-   ##
=======================================
  Coverage   92.86%   92.86%           
=======================================
  Files         223      223           
  Lines       17436    17443    +7     
=======================================
+ Hits        16192    16199    +7     
  Misses       1244     1244           
```


| [Impacted Files](https://codecov.io/gh/OpenSourceEconomics/estimagic/pull/411?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=OpenSourceEconomics) | Coverage Δ | |
|---|---|---|
| [tests/optimization/test\_tiktak.py](https://codecov.io/gh/OpenSourceEconomics/estimagic/pull/411/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=OpenSourceEconomics#diff-dGVzdHMvb3B0aW1pemF0aW9uL3Rlc3RfdGlrdGFrLnB5) | `93.25% <ø> (ø)` | |
| [src/estimagic/optimization/tiktak.py](https://codecov.io/gh/OpenSourceEconomics/estimagic/pull/411/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=OpenSourceEconomics#diff-c3JjL2VzdGltYWdpYy9vcHRpbWl6YXRpb24vdGlrdGFrLnB5) | `92.63% <100.00%> (+0.18%)` | :arrow_up: |
| [tests/optimization/test\_multistart.py](https://codecov.io/gh/OpenSourceEconomics/estimagic/pull/411/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=OpenSourceEconomics#diff-dGVzdHMvb3B0aW1pemF0aW9uL3Rlc3RfbXVsdGlzdGFydC5weQ==) | `100.00% <100.00%> (ø)` | |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=OpenSourceEconomics). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=OpenSourceEconomics)


Patch:
diff --git a/src/estimagic/optimization/tiktak.py b/src/estimagic/optimization/tiktak.py
index ea21eb5..31f511f 100644
--- a/src/estimagic/optimization/tiktak.py
+++ b/src/estimagic/optimization/tiktak.py
@@ -170,6 +170,7 @@ def run_multistart_optimization(
             starts=starts,
             results=batch_results,
             convergence_criteria=convergence_criteria,
+            primary_key=primary_key,
         )
         opt_counter += len(batch)
         scheduled_steps = scheduled_steps[len(batch) :]
@@ -427,7 +428,9 @@ def get_batched_optimization_sample(sorted_sample, n_optimizations, batch_size):
     return batched
 
 
-def update_convergence_state(current_state, starts, results, convergence_criteria):
+def update_convergence_state(
+    current_state, starts, results, convergence_criteria, primary_key
+):
     """Update the state of all quantities related to convergence.
 
     Args:
@@ -442,6 +445,8 @@ def update_convergence_state(current_state, starts, results, convergence_criteri
         starts (list): List of starting points for local optimizations.
         results (list): List of results from local optimizations.
         convergence_criteria (dict): Dict with the entries "xtol" and "max_discoveries"
+        primary_key: The primary criterion entry of the local optimizer. Needed to
+            interpret the output of the internal criterion function.
 
 
     Returns:
@@ -456,17 +461,38 @@ def update_convergence_state(current_state, starts, results, convergence_criteri
     best_y = current_state["best_y"]
     best_res = current_state["best_res"]
 
+    # get indices of local optimizations that did not fail
     valid_indices = [i for i, res in enumerate(results) if not isinstance(res, str)]
 
+    # If all local optimizations failed, return early so we don't have to worry about
+    # index errors later.
     if not valid_indices:
         return current_state, False
 
+    # ==================================================================================
+    # reduce eveything to valid optimizations
+    # ==================================================================================
     valid_results = [results[i] for i in valid_indices]
     valid_starts = [starts[i] for i in valid_indices]
-
     valid_new_x = [res["solution_x"] for res in valid_results]
-    valid_new_y = [res["solution_criterion"] for res in valid_results]
+    valid_new_y = []
+
+    # make the criterion output scalar if a least squares optimizer returns an
+    # array as solution_criterion.
+    for res in valid_results:
+        if np.isscalar(res["solution_criterion"]):
+            valid_new_y.append(res["solution_criterion"])
+        else:
+            valid_new_y.append(
+                aggregate_func_output_to_value(
+                    f_eval=res["solution_criterion"],
+                    primary_key=primary_key,
+                )
+            )
 
+    # ==================================================================================
+    # accept new best point if we find a new lowest function value
+    # ==================================================================================
     best_index = np.argmin(valid_new_y)
     if valid_new_y[best_index] <= best_y:
         best_x = valid_new_x[best_index]
@@ -478,6 +504,9 @@ def update_convergence_state(current_state, starts, results, convergence_criteri
     elif best_res is None:
         best_res = valid_results[best_index]
 
+    # ==================================================================================
+    # update history and state
+    # ==================================================================================
     new_x_history = current_state["x_history"] + valid_new_x
     all_x = np.array(new_x_history)
     relative_diffs = (all_x - best_x) / np.clip(best_x, 0.1, np.inf)


Test patch:
diff --git a/tests/optimization/test_multistart.py b/tests/optimization/test_multistart.py
index f800755..51e17df 100644
--- a/tests/optimization/test_multistart.py
+++ b/tests/optimization/test_multistart.py
@@ -239,3 +239,17 @@ def test_with_ackley():
             "convergence_max_discoveries": 10,
         },
     )
+
+
+def test_multistart_with_least_squares_optimizers():
+    est = minimize(
+        criterion=sos_dict_criterion,
+        params=np.array([-1, 1.0]),
+        lower_bounds=np.full(2, -10.0),
+        upper_bounds=np.full(2, 10.0),
+        algorithm="scipy_ls_trf",
+        multistart=True,
+        multistart_options={"n_samples": 3, "share_optimizations": 1.0},
+    )
+
+    aaae(est.params, np.zeros(2))
diff --git a/tests/optimization/test_tiktak.py b/tests/optimization/test_tiktak.py
index d1bcad7..7445ee6 100644
--- a/tests/optimization/test_tiktak.py
+++ b/tests/optimization/test_tiktak.py
@@ -153,6 +153,7 @@ def test_update_state_converged(current_state, starts, results):
         starts=starts,
         results=results,
         convergence_criteria=criteria,
+        primary_key="value",
     )
 
     aaae(new_state["best_x"], np.arange(3))
@@ -176,6 +177,7 @@ def test_update_state_not_converged(current_state, starts, results):
         starts=starts,
         results=results,
         convergence_criteria=criteria,
+        primary_key="value",
     )
 
     assert not is_converged