Correct the order of rvs sent to compile_dlogp in find_MAP (#5923)

Author: quantheory

pymc/tests/test_starting.py

@@ -98,6 +98,32 @@ def test_find_MAP():
     close_to(map_est2["sigma"], 1, tol)
 
 
+def test_find_MAP_issue_5923():
+    # Test that gradient-based minimization works well regardless of the order
+    # of variables in `vars`, and even when starting a reasonable distance from
+    # the MAP.
+    tol = 2.0**-11  # 16 bit machine epsilon, a low bar
+    data = np.random.randn(100)
+    # data should be roughly mean 0, std 1, but let's
+    # normalize anyway to get it really close
+    data = (data - np.mean(data)) / np.std(data)
+
+    with Model():
+        mu = Uniform("mu", -1, 1)
+        sigma = Uniform("sigma", 0.5, 1.5)
+        Normal("y", mu=mu, tau=sigma**-2, observed=data)
+
+        start = {"mu": -0.5, "sigma": 1.25}
+        map_est1 = starting.find_MAP(progressbar=False, vars=[mu, sigma], start=start)
+        map_est2 = starting.find_MAP(progressbar=False, vars=[sigma, mu], start=start)
+
+    close_to(map_est1["mu"], 0, tol)
+    close_to(map_est1["sigma"], 1, tol)
+
+    close_to(map_est2["mu"], 0, tol)
+    close_to(map_est2["sigma"], 1, tol)
+
+
 def test_find_MAP_issue_4488():
     # Test for https://github.com/pymc-devs/pymc/issues/4488
     with Model() as m:

pymc/tuning/starting.py

@@ -110,17 +110,17 @@ def find_MAP(
     start = ipfn(seed)
     model.check_start_vals(start)
 
-    var_names = {var.name for var in vars}
+    vars_dict = {var.name: var for var in vars}
     x0 = DictToArrayBijection.map(
-        {var_name: value for var_name, value in start.items() if var_name in var_names}
+        {var_name: value for var_name, value in start.items() if var_name in vars_dict}
     )
 
     # TODO: If the mapping is fixed, we can simply create graphs for the
     # mapping and avoid all this bijection overhead
     compiled_logp_func = DictToArrayBijection.mapf(model.compile_logp(jacobian=False), start)
     logp_func = lambda x: compiled_logp_func(RaveledVars(x, x0.point_map_info))
 
-    rvs = [model.values_to_rvs[value] for value in vars]
+    rvs = [model.values_to_rvs[vars_dict[name]] for name, _, _ in x0.point_map_info]
     try:
         # This might be needed for calls to `dlogp_func`
         # start_map_info = tuple((v.name, v.shape, v.dtype) for v in vars)