Merge pull request #118 from leila-pujal/derivatives

Speed up gradient and hessian computation Closes #99

Author: leila-pujal

gbasis/evals/_deriv.py

@@ -140,3 +140,238 @@ def _eval_deriv_contractions(coords, orders, center, angmom_comps, alphas, prim_
 
     norm = norm.T[:, :, np.newaxis]
     return np.tensordot(prim_coeffs, norm * zeroth_part * deriv_part, (0, 0))
+
+
+def _eval_first_second_order_deriv_contractions(
+        coords, orders, center, angmom_comps, alphas, prim_coeffs, norm
+):
+    """Return the evaluation of direct 1st and 2nd derivative orders of a Cartesian contraction.
+
+    Parameters
+    ----------
+    coords : np.ndarray(N, 3)
+        Point in space where the derivative of the Gaussian primitive is evaluated.
+        Coordinates must be given as a two dimensional array, even if only one point is given.
+    orders : np.ndarray(3,)
+        Orders of the derivative.
+        All orders must be lower than 2.
+        Negative orders are treated as zero orders.
+    center : np.ndarray(3,)
+        Center of the Gaussian primitive.
+    angmom_comps : np.ndarray(L, 3)
+        Components of the angular momentum, :math:`(a_x, a_y, a_z)`.
+        Angular momentum components must be given as a two dimensional array, even if only one
+        set of components is given.
+    alphas : np.ndarray(K,)
+        Values of the (square root of the) precisions of the primitives.
+    prim_coeffs : np.ndarray(K, M)
+        Contraction coefficients of the primitives.
+        The coefficients always correspond to generalized contractions, i.e. two-dimensional array
+        where the first index corresponds to the primitive and the second index corresponds to the
+        contraction (with the same exponents and angular momentum).
+    norm : np.ndarray(L, K)
+        Normalization constants for the primitives in each contraction.
+
+    Returns
+    -------
+    derivative : np.ndarray(M, L, N)
+        Evaluation of the derivative at each given coordinate.
+        Dimension 0 corresponds to the contraction, with `M` as the number of given contractions.
+        Dimension 1 corresponds to the angular momentum vector, ordered as in `angmom_comps`.
+        Dimension 2 corresponds to the point at which the derivative is evaluated, ordered as in
+        `coords`.
+
+    Notes
+    -----
+    The input is not checked. This means that you must provide the parameters as they are specified
+    in the docstring. They must all be `numpy` arrays with the **correct shape**.
+
+    Pople style basis sets are not supported. If multiple angular momentum vectors (with different
+    angular momentum) and multiple contraction coefficients are provided, it is **not assumed** that
+    the angular momentum vector should be paired up with the contraction coefficients. In fact, each
+    angular momentum vector will create multiple contractions according to the given coefficients.
+
+    """
+
+    # Useful variables
+    new_coords = coords.T - center[None, :].T
+    gauss = np.exp(-alphas[:, None, None] * (new_coords ** 2))
+
+    # Filters derivative orders
+    indices_noderiv = orders <= 0
+    indices_first_deriv = orders == 1
+    indices_second_deriv = orders == 2
+
+    # Zeroth order derivative
+    raw_zeroth_coords = new_coords[indices_noderiv]
+    zeroth_angmom_comps = angmom_comps.T[indices_noderiv]
+    zeroth_gauss = gauss[:, indices_noderiv]
+    zeroth_coords = raw_zeroth_coords[:, None, :] ** zeroth_angmom_comps[:, :, None]
+    raw_zeroth_deriv = zeroth_coords[None, :, :, :] * zeroth_gauss[:, :, None, :]
+    zeroth_deriv = np.prod(raw_zeroth_deriv, axis=1)
+
+    # Initialize first and second derivative variables with shape = zeroth_deriv.shape
+    first_deriv = np.ones(zeroth_deriv.shape)
+    second_deriv = np.ones(zeroth_deriv.shape)
+
+    # Calling 1st and 2nd derivatives functions for different combination of orders
+    if indices_first_deriv.any():
+        first_deriv = _first_derivative(new_coords, gauss, indices_first_deriv,
+                                        angmom_comps, alphas)
+        if indices_second_deriv.any():
+            second_deriv = _second_derivative(new_coords, gauss, indices_second_deriv,
+                                              angmom_comps, alphas)
+    elif indices_second_deriv.any():
+        second_deriv = _second_derivative(new_coords, gauss, indices_second_deriv,
+                                          angmom_comps, alphas)
+        if indices_first_deriv.any():
+            first_deriv = _first_derivative(new_coords, gauss, indices_first_deriv,
+                                            angmom_comps, alphas)
+    # Combining all the derivatives
+    norm = norm.T[:, :, np.newaxis]
+    output = np.tensordot(prim_coeffs, norm * zeroth_deriv * first_deriv * second_deriv, (0, 0))
+
+    return output
+
+
+def _first_derivative(center_coords, gauss, indices_first_deriv, angmom_comps, alphas):
+    r"""Help function for calculation of explicit first derivative order for contracted gaussian.
+
+    Parameters
+    ----------
+    center_coords : np.ndarray(N, 3)
+        Shifted coordinates center_coords = coords - center
+    gauss : np.ndarray(K, 3, N)
+        variable containing ..math::e^{-\alpha\left(x-X_{A}\right)^{2}}
+    indices_first_deriv: boolean np.array(L, 3)
+        array contaning boolean values corresponding to coordinates for first derivative
+    angmom_comps : np.ndarray(L, 3)
+        Components of the angular momentum, :math:`(a_x, a_y, a_z)`.
+        Angular momentum components must be given as a two dimensional array, even if only one
+        set of components is given.
+    alphas : np.ndarray(K,)
+        Values of the (square root of the) precisions of the primitives.
+
+    Returns
+    -------
+    first order derivative : np.ndarray(K, L, N)
+        Evaluation of first derivative at each given coordinate.
+        Dimension 0 corresponds to number of primitives.
+        Dimension 1 corresponds to the angular momentum vector, ordered as in `angmom_comps`.
+        Dimension 2 corresponds to the point at which the derivative is evaluated, ordered as in
+        `coords`.
+
+    Notes
+    -----
+    This is a helper function for _eval_first_second_order_deriv_contractions() and gives an
+    intermediate output that needs to be further processed to match general organization in Gbasis
+
+    """
+
+    first_coords = center_coords[indices_first_deriv]
+    first_gauss = gauss[:, indices_first_deriv]
+    first_ang_comp = angmom_comps.T[indices_first_deriv]
+    # Indices to filter for ang momentum at the end
+    n_0_indices = first_ang_comp == 0
+
+    power_part_1 = first_ang_comp - 1
+    # NOTE: the negative indices must be turned into zeros because (x-Xa) terms are sometimes
+    # zero (and negative power is undefined).
+    power_part_1[power_part_1 < 0] = 0
+    part1 = first_coords[:, None, :] ** power_part_1[:, :, None]
+    part2 = (2 * alphas[:, None, None]) * (first_coords ** 2)
+    part2 = first_ang_comp[None, :, :, None] - part2[:, :, None, :]
+    # NOTE: Using an array of ones with same shape as first_ang_comp to power part2_zero_ang_mom
+    # variable in order to get the same shape as part2. This is done in order to make easier
+    # to filter at the end for the angular components corresponding to n=0
+    array_ones = np.ones(first_ang_comp.shape)
+    part2_n_0 = -2 * alphas[:, None, None, None] * (
+            first_coords[:, None, :] ** array_ones[:, :, None])
+    raw_first_deriv = part1 * part2
+    # Substitute angular components n=0 with correct derivative
+    raw_first_deriv[:, n_0_indices, :] = part2_n_0[:, n_0_indices, :]
+    raw_first_deriv = raw_first_deriv * first_gauss[:, :, None, :]
+    first_deriv = np.prod(raw_first_deriv, axis=1)
+
+    return first_deriv
+
+
+def _second_derivative(center_coords, gauss, indices_second_deriv, angmom_comps, alphas):
+    r"""Help function for calculation of explicit second derivative order for contracted gaussian.
+
+    Parameters
+    ----------
+    center_coords : np.ndarray(N, 3)
+        Shifted coordinates center_coords = coords - center
+    gauss : np.ndarray(K, 3, N)
+        variable containing ..math::e^{-\alpha\left(x-X_{A}\right)^{2}}
+    indices_second_deriv: boolean np.array(L, 3)
+        array contaning boolean values corresponding to coordinates for second derivative
+    angmom_comps : np.ndarray(L, 3)
+        Components of the angular momentum, :math:`(a_x, a_y, a_z)`.
+        Angular momentum components must be given as a two dimensional array, even if only one
+        set of components is given.
+    alphas : np.ndarray(K,)
+        Values of the (square root of the) precisions of the primitives.
+
+    Returns
+    -------
+    first order derivative : np.ndarray(K, L, N)
+        Evaluation of second derivative at each given coordinate.
+        Dimension 0 corresponds to number of primitives.
+        Dimension 1 corresponds to the angular momentum vector, ordered as in `angmom_comps`.
+        Dimension 2 corresponds to the point at which the derivative is evaluated, ordered as in
+        `coords`.
+
+    Notes
+    -----
+    This is a helper function for _eval_first_second_order_deriv_contractions() and gives an
+    intermediate output that needs to be further processed to match general organization in Gbasis
+
+    """
+
+    second_coords = center_coords[indices_second_deriv]
+    second_gauss = gauss[:, indices_second_deriv]
+    second_ang_comp = angmom_comps.T[indices_second_deriv]
+    # NOTE: As for the first derivative, using an array of ones with shape=second_ang_comp.shape
+    # to match the shape of different variables, corresponding to calculations of n=0 and n=1
+    # second derivatives, to make easier at the end to combine them.
+    array_ones = np.ones(second_ang_comp.shape)
+    # Indices to filter for ang momentum at the end
+    n_1_indices = second_ang_comp == 1
+    n_2_indices = second_ang_comp >= 2
+
+    # angular momentum == 0
+    total_n_0 = ((4 * alphas[:, None, None] ** 2) * (second_coords ** 2)) - \
+                (2 * alphas[:, None, None])
+    raw_second_deriv = total_n_0[:, :, None, :] ** array_ones[None, :, :, None]
+    # angular momentum == 1
+    if any(second_ang_comp[0] == 1):
+        total_n_1 = ((4 * alphas[:, None, None] ** 2) * (second_coords ** 3)) \
+                    - ((6 * alphas[:, None, None]) * second_coords)
+        total_n_1 = total_n_1[:, :, None, :] ** array_ones[None, :, :, None]
+        # Substitute angular components n=1 with correct derivative
+        raw_second_deriv[:, n_1_indices, :] = total_n_1[:, n_1_indices, :]
+        # angular momentum >= 2
+        if any(second_ang_comp[0] >= 2):
+            # Calculating ..math:: \left(x-X_{A}\right)^{n-2}
+            power_part_1 = second_ang_comp - 2
+            # NOTE: the negative indices must be turned into zeros because (x-Xa)
+            # terms are sometimes zeros (and negative power is undefined).
+            power_part_1[power_part_1 < 0] = 0
+            part1_n_2 = second_coords[:, None, :] ** (power_part_1[:, :, None])
+            # Calculating
+            # ..math:: 4 \alpha^{2}\left(x-X_{A}\right)^{4}-
+            # \alpha(4 n+2)\left(x-X_{A}\right)^{2}+n(n-1)
+            part2_1_n_2 = (4 * alphas[:, None, None] ** 2) * (second_coords ** 4)
+            part2_2_n_2 = alphas[:, None, None, None] * (4 * second_ang_comp[:, :, None] + 2) \
+                                                      * second_coords[:, None, :] ** 2
+            part2_3_n_2 = second_ang_comp * (second_ang_comp - 1)
+            part2_n_2 = part2_1_n_2[:, :, None, :] - part2_2_n_2 + part2_3_n_2[None, :, :, None]
+            total_n_2 = part1_n_2[None, :, :, :] * part2_n_2
+            # Substitute angular components n=2 with correct derivative
+            raw_second_deriv[:, n_2_indices, :] = total_n_2[:, n_2_indices, :]
+    raw_second_deriv = raw_second_deriv * second_gauss[:, :, None, :]
+    second_deriv = np.prod(raw_second_deriv, axis=1)
+
+    return second_deriv