BUG: Fix problem with Scipy.integrate.solve_bvp for big problems (sci…

…py#12556)

* BUG: Fix problem with Scipy.integrate.solve_bvp for big problems with parameters.

This commit solves fatal errors when solve_bvp is used for a big problem, i.e. size > 50,
and the unknowns include parameters that need solving.
This commit includes tests, with and without Jacobians for both the function and the boundary conditions.
The new tests adapt already existing tests to make them easier to compare and understand.

THANKS.txt

@@ -240,6 +240,7 @@ Luigi F. Cruz for adding time/frequency domain option to scipy.signal.resample.
 Wesley Alves for improvements to scipy.stats.jarque_bera and scipy.stats.shapiro
 Mark Borgerding for contributing linalg.convolution_matrix.
 Shashaank N for contributions to scipy.signal.
+Frank Torres for fixing a bug with solve_bvp for large problems.
 
 Institutions
 ------------

scipy/integrate/_bvp.py

@@ -143,12 +143,11 @@ def compute_jac_indices(n, m, k):
 def stacked_matmul(a, b):
     """Stacked matrix multiply: out[i,:,:] = np.dot(a[i,:,:], b[i,:,:]).
 
-    In our case, a[i, :, :] and b[i, :, :] are always square.
+    Empirical optimization. Use outer Python loop and BLAS for large
+    matrices, otherwise use a single einsum call.
     """
-    # Empirical optimization. Use outer Python loop and BLAS for large
-    # matrices, otherwise use a single einsum call.
     if a.shape[1] > 50:
-        out = np.empty_like(a)
+        out = np.empty((a.shape[0], a.shape[1], b.shape[2]))
         for i in range(a.shape[0]):
             out[i] = np.dot(a[i], b[i])
         return out

scipy/integrate/tests/test_bvp.py

@@ -152,6 +152,76 @@ def big_sol(x, n):
     return x
 
 
+def big_fun_with_parameters(x, y, p):
+    """ Big version of sl_fun, with two parameters.
+
+    The two differential equations represented by sl_fun are broadcast to the
+    number of rows of y, rotating between the parameters p[0] and p[1].
+    Here are the differential equations:
+
+        dy[0]/dt = y[1]
+        dy[1]/dt = -p[0]**2 * y[0]
+        dy[2]/dt = y[3]
+        dy[3]/dt = -p[1]**2 * y[2]
+        dy[4]/dt = y[5]
+        dy[5]/dt = -p[0]**2 * y[4]
+        dy[6]/dt = y[7]
+        dy[7]/dt = -p[1]**2 * y[6]
+        .
+        .
+        .
+
+    """
+    f = np.zeros_like(y)
+    f[::2] = y[1::2]
+    f[1::4] = -p[0]**2 * y[::4]
+    f[3::4] = -p[1]**2 * y[2::4]
+    return f
+
+
+def big_fun_with_parameters_jac(x, y, p):
+    # big version of sl_fun_jac, with two parameters
+    n, m = y.shape
+    df_dy = np.zeros((n, n, m))
+    df_dy[range(0, n, 2), range(1, n, 2)] = 1
+    df_dy[range(1, n, 4), range(0, n, 4)] = -p[0]**2
+    df_dy[range(3, n, 4), range(2, n, 4)] = -p[1]**2
+
+    df_dp = np.zeros((n, 2, m))
+    df_dp[range(1, n, 4), 0] = -2 * p[0] * y[range(0, n, 4)]
+    df_dp[range(3, n, 4), 1] = -2 * p[1] * y[range(2, n, 4)]
+
+    return df_dy, df_dp
+
+
+def big_bc_with_parameters(ya, yb, p):
+    # big version of sl_bc, with two parameters
+    return np.hstack((ya[::2], yb[::2], ya[1] - p[0], ya[3] - p[1]))
+
+
+def big_bc_with_parameters_jac(ya, yb, p):
+    # big version of sl_bc_jac, with two parameters
+    n = ya.shape[0]
+    dbc_dya = np.zeros((n + 2, n))
+    dbc_dyb = np.zeros((n + 2, n))
+
+    dbc_dya[range(n // 2), range(0, n, 2)] = 1
+    dbc_dyb[range(n // 2, n), range(0, n, 2)] = 1
+
+    dbc_dp = np.zeros((n + 2, 2))
+    dbc_dp[n, 0] = -1
+    dbc_dya[n, 1] = 1
+    dbc_dp[n + 1, 1] = -1
+    dbc_dya[n + 1, 3] = 1
+
+    return dbc_dya, dbc_dyb, dbc_dp
+
+
+def big_sol_with_parameters(x, p):
+    # big version of sl_sol, with two parameters
+    return np.vstack((np.sin(p[0] * x), np.sin(p[1] * x)))
+
+
 def shock_fun(x, y):
     eps = 1e-3
     return np.vstack((
@@ -532,6 +602,43 @@ def test_big_problem():
     assert_allclose(sol.sol(sol.x, 1), sol.yp, rtol=1e-10, atol=1e-10)
 
 
+def test_big_problem_with_parameters():
+    n = 30
+    x = np.linspace(0, np.pi, 5)
+    x_test = np.linspace(0, np.pi, 100)
+    y = np.ones((2 * n, x.size))
+
+    for fun_jac in [None, big_fun_with_parameters_jac]:
+        for bc_jac in [None, big_bc_with_parameters_jac]:
+            sol = solve_bvp(big_fun_with_parameters, big_bc_with_parameters, x,
+                            y, p=[0.5, 0.5], fun_jac=fun_jac, bc_jac=bc_jac)
+
+            assert_equal(sol.status, 0)
+            assert_(sol.success)
+
+            assert_allclose(sol.p, [1, 1], rtol=1e-4)
+
+            sol_test = sol.sol(x_test)
+
+            for isol in range(0, n, 4):
+                assert_allclose(sol_test[isol],
+                                big_sol_with_parameters(x_test, [1, 1])[0],
+                                rtol=1e-4, atol=1e-4)
+                assert_allclose(sol_test[isol + 2],
+                                big_sol_with_parameters(x_test, [1, 1])[1],
+                                rtol=1e-4, atol=1e-4)
+
+            f_test = big_fun_with_parameters(x_test, sol_test, [1, 1])
+            r = sol.sol(x_test, 1) - f_test
+            rel_res = r / (1 + np.abs(f_test))
+            norm_res = np.sum(rel_res ** 2, axis=0) ** 0.5
+            assert_(np.all(norm_res < 1e-3))
+
+            assert_(np.all(sol.rms_residuals < 1e-3))
+            assert_allclose(sol.sol(sol.x), sol.y, rtol=1e-10, atol=1e-10)
+            assert_allclose(sol.sol(sol.x, 1), sol.yp, rtol=1e-10, atol=1e-10)
+
+
 def test_shock_layer():
     x = np.linspace(-1, 1, 5)
     x_test = np.linspace(-1, 1, 100)