Merge df2259c into 1dc72e8

smt/applications/mfk.py

@@ -954,7 +954,8 @@ def _predict_derivatives(self, x, kx):
             # scaled predictor
             dy_dx[:, i] = np.ravel(df_dx - 2 * theta[kx] * np.dot(d_dx * r_, gamma))
 
-        return dy_dx[:, -1] * self.y_std / self.X_scale[kx]
+        deriv = dy_dx[:, -1] * self.y_std / self.X_scale[kx]
+        return deriv.reshape((x.shape[0], self.ny))
 
     def _get_theta(self, i):
         return self.optimal_theta[i]

smt/surrogate_models/krg_based.py

@@ -1643,18 +1643,43 @@ def _predict_variances(self, x: np.ndarray, is_acting=None) -> np.ndarray:
 
     def _predict_variance_derivatives(self, x, kx):
         """
-        Provide the derivative of the variance of the model at a set of points
+        Provide the derivatives of the variance of the model at a set of points
+
         Parameters
         -----------
         x : np.ndarray [n_evals, dim]
             Evaluation point input variable values
         kx : int
             The 0-based index of the input variable with respect to which derivatives are desired.
 
+        Returns
+        -------
+        derived_variance:  np.ndarray
+            The derivatives wrt kx-th component of the variance of the kriging model
+        """
+        return self._internal_predict_variance(x, kx)
+
+    def _predict_variance_gradient(self, x):
+        """
+        Provide the gradient of the variance of the model at a given point
+        (ie the derivatives wrt to all component at a unique point x)
+
+        Parameters
+        -----------
+        x : np.ndarray [1, dim]
+            Evaluation point input variable values
+
         Returns
         -------
          derived_variance:  np.ndarray
-             The jacobian of the variance of the kriging model (the kx-th derivative)
+            The gradient of the variance of the kriging model
+        """
+        return self._internal_predict_variance(x)
+
+    def _internal_predict_variance(self, x, kx=None):
+        """
+        When kx is None gradient is computed at the location x
+        otherwise partial derivatives wrt kx-th component of a set of points x
         """
 
         # Initialization
@@ -1676,11 +1701,14 @@ def _predict_variance_derivatives(self, x, kx):
         r, dr = self._correlation_types[self.options["corr"]](
             theta, d, derivative_params=derivative_dic
         )
-        r = r.reshape(n_eval, self.nt)
-        drx = dr[:, kx]
-        drx = drx.reshape(n_eval, self.nt)
+        if kx is None:
+            rt = linalg.solve_triangular(C, r, lower=True)
+            drx = dr.T
+        else:
+            r = r.reshape(n_eval, self.nt)
+            rt = linalg.solve_triangular(C, r.T, lower=True)
+            drx = dr[:, kx].reshape(n_eval, self.nt)
 
-        rt = linalg.solve_triangular(C, r.T, lower=True)
         invKr = linalg.solve_triangular(C.T, rt)
         p1 = 2 * np.dot(drx, invKr).T
 
@@ -1690,7 +1718,10 @@ def _predict_variance_derivatives(self, x, kx):
         rho2 = linalg.solve_triangular(C, F, lower=True)
         invKF = linalg.solve_triangular(C.T, rho2)
 
-        A = f_x.T - np.dot(r, invKF)
+        if kx is None:
+            A = f_x.T - np.dot(r.T, invKF)
+        else:
+            A = f_x.T - np.dot(r, invKF)
 
         B = np.dot(F.T, invKF)
         rho3 = linalg.cholesky(B, lower=True)
@@ -1708,17 +1739,27 @@ def _predict_variance_derivatives(self, x, kx):
                 + "universal kriging using a linear trend"
             )
 
-        dA = df[:, kx].T - np.dot(drx, invKF)
+        if kx is None:
+            dA = df.T - np.dot(dr.T, invKF)
+        else:
+            dA = df[:, kx].T - np.dot(drx, invKF)
+
         p3 = 2 * np.dot(dA, D).T
 
         # prime : derivative of MSE
         # MSE ~1.0 - (rt**2.0).sum(axis=0) + (u**2.0).sum(axis=0)
         prime = 0 - p1 + p3
         ## scaling factors
-        x_std = self.X_scale[kx]
-        derived_variance = np.array((np.outer(sigma2, np.diag(prime.T)) / x_std))
-
-        return np.atleast_2d(derived_variance.T)
+        if kx is None:
+            derived_variance = []
+            x_std = np.resize(self.X_scale, self.nx)
+            for i in range(len(x_std)):
+                derived_variance.append(sigma2 * prime.T[i] / x_std[i])
+            return np.array(derived_variance).T
+        else:
+            x_std = self.X_scale[kx]
+            derived_variance = np.array((np.outer(sigma2, np.diag(prime.T)) / x_std))
+            return np.atleast_2d(derived_variance).T
 
     def _optimize_hyperparam(self, D):
         """

smt/surrogate_models/qp.py

@@ -160,4 +160,5 @@ def _predict_values(self, x):
         M = self._response_surface(x)
         y_ = np.dot(M, self.coef)
         y = (self.y_mean + self.y_std * y_).ravel()
+        y = y.reshape((x.shape[0], self.ny))
         return y