diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index e674b658..0072ab29 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -58,7 +58,6 @@ jobs: # inspired by https://github.com/orgs/community/discussions/26724 # only push to gh-pages if the master branch has been updated - name: GitHub Pages Action - # if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }} uses: peaceiris/actions-gh-pages@v3 with: github_token: ${{ secrets.GITHUB_TOKEN }} diff --git a/gbasis/autograd/playground.py b/gbasis/autograd/playground.py new file mode 100644 index 00000000..55457d8b --- /dev/null +++ b/gbasis/autograd/playground.py @@ -0,0 +1,115 @@ +from autograd import grad +import autograd.numpy as np +from functools import partial + +def eval_contractions(coords, center, angmom_comps, alphas, prim_coeffs, norm): + """Return the evaluation 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. + 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. + + """ + # pylint: disable=R0914 + # NOTE: following convention will be used to organize the axis of the multidimensional arrays + # axis 0 = index for term in hermite polynomial (size: min(K, n)) where n is the order in given + # dimension + # axis 1 = index for primitive (size: K) + # axis 2 = index for dimension (x, y, z) of coordinate (size: 3) + # axis 3 = index for angular momentum vector (size: L) + # axis 4 = index for coordinate (out of a grid) (size: N) + # adjust the axis + coords = coords.T[np.newaxis, np.newaxis, :, np.newaxis, :] + # NOTE: if `coord` is two dimensional (3, N), then coords has shape (1, 1, 3, 1, N). If it is + # one dimensional (3,), then coords has shape (1, 1, 3, 1) + # NOTE: `order` is still assumed to be a one dimensional + center = center[np.newaxis, np.newaxis, :, np.newaxis, np.newaxis] + angmom_comps = angmom_comps.T[np.newaxis, np.newaxis, :, :, np.newaxis] + # NOTE: if `angmom_comps` is two-dimensional (3, L), has shape (1, 1, 3, L). If it is one + # dimensional (3, ) then it has shape (1, 1, 3) + alphas = alphas[np.newaxis, :, np.newaxis, np.newaxis, np.newaxis] + # NOTE: `prim_coeffs` will be used as a 1D array + + # shift coordinates + coords = coords - center + # useful variables + gauss = np.exp(-alphas * coords ** 2) + + zeroth_part = np.prod(coords ** angmom_comps * gauss, axis=(0, 2)) + # NOTE: `zeroth_part` now has axis 0 for primitives, axis 1 for angular momentum vector, and + # axis 2 for coordinate + + norm = norm.T[:, :, np.newaxis] + return np.tensordot(prim_coeffs, norm * zeroth_part, (0, 0)).flatten() + +coords = np.random.rand(1, 1) +center = np.random.rand(1) +angmom_comps = np.random.randint(0, 2, (1, 1)) +alphas = np.random.rand(1) +prim_coeffs = np.random.rand(1, 1) +norm = np.ones((1, 1)) + +output = eval_contractions(coords, center, angmom_comps, alphas, prim_coeffs, norm) +print(output) +eval_contruction_merged = partial(eval_contractions, + center=center, + angmom_comps=angmom_comps, + alphas=alphas, + prim_coeffs=prim_coeffs, + norm=norm) +eval_contractions_grad = grad(eval_contruction_merged) +print(eval_contractions_grad(coords)) + +# coords = coords.T[np.newaxis, np.newaxis, :, np.newaxis, :] +# center = center[np.newaxis, np.newaxis, :, np.newaxis, np.newaxis] +# angmom_comps = angmom_comps.T[np.newaxis, np.newaxis, :, :, np.newaxis] +# alphas = alphas[np.newaxis, :, np.newaxis, np.newaxis, np.newaxis] +# coords = coords - center +# gauss = np.exp(-alphas * coords ** 2) + +# zeroth_part = np.prod(coords ** angmom_comps * gauss, axis=(0, 2)) +# norm = norm.T[:, :, np.newaxis] +# np.tensordot(prim_coeffs, norm * zeroth_part, (0, 0)) + +# print() +# print(eval_contractions_grad(coords).shape) + diff --git a/gbasis/gradient/__init__.py b/gbasis/gradient/__init__.py new file mode 100644 index 00000000..6bb5e4e7 --- /dev/null +++ b/gbasis/gradient/__init__.py @@ -0,0 +1 @@ +"""Module for differentiating Gaussian functions using automatic differention.""" \ No newline at end of file diff --git a/gbasis/gradient/ad.py b/gbasis/gradient/ad.py new file mode 100644 index 00000000..beccd59f --- /dev/null +++ b/gbasis/gradient/ad.py @@ -0,0 +1,197 @@ +from autograd import grad +import autograd.numpy as np + +# it looks like there is a plan to do: +# 1. rewrite the calculation of normalization constants (see if there is a dependency on the R_a in the normalization constant) +# 2. rewrite the calculation of contractions +# 3. use autograd to calculate the gradient of the contractions + +def eval_contractions(coords, R_x, R_y, R_z, angmom_comps, alphas, prim_coeffs, norm): + """Return the evaluation of a Cartesian contraction. + Note that it was addapted from the evals/_deriv.py file in the gbasis package to match the autograd API. + + Parameters + ---------- + coords : np.ndarray(3) + Point in space where the derivative of the Gaussian primitive is evaluated. + Coordinates must be given as a one(!) dimensional array. + center : np.ndarray(3,) + Center of the Gaussian primitive. + angmom_comps : np.ndarray(1, 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, 1) + 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(1, 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. + + """ + # pylint: disable=R0914 + # NOTE: following convention will be used to organize the axis of the multidimensional arrays + # axis 0 = index for term in hermite polynomial (size: min(K, n)) where n is the order in given + # dimension + # axis 1 = index for primitive (size: K) + # axis 2 = index for dimension (x, y, z) of coordinate (size: 3) + # axis 3 = index for angular momentum vector (size: L) + # axis 4 = index for coordinate (out of a grid) (size: N) + # adjust the axis + + coords = coords.reshape(1, 3) # this is done so the following code from the gbasis package works + coords = coords.T[np.newaxis, np.newaxis, :, np.newaxis, :] + # NOTE: if `coord` is two dimensional (3, N), then coords has shape (1, 1, 3, 1, N). If it is + # one dimensional (3,), then coords has shape (1, 1, 3, 1) + # NOTE: `order` is still assumed to be a one dimensional + center = np.array([R_x, R_y, R_z]) + center = center[np.newaxis, np.newaxis, :, np.newaxis, np.newaxis] + + angmom_comps = angmom_comps.reshape(1, 3) + angmom_comps = angmom_comps.T[np.newaxis, np.newaxis, :, :, np.newaxis] #this line + # NOTE: if `angmom_comps` is two-dimensional (3, L), has shape (1, 1, 3, L). If it is one + # dimensional (3, ) then it has shape (1, 1, 3) + alphas = alphas[np.newaxis, :, np.newaxis, np.newaxis, np.newaxis] + # NOTE: `prim_coeffs` will be used as a 1D array + + # shift coordinates + coords = coords - center + # useful variables + gauss = np.exp(-alphas * coords ** 2) + + zeroth_part = np.prod(coords ** angmom_comps * gauss, axis=(0, 2)) + # NOTE: `zeroth_part` now has axis 0 for primitives, axis 1 for angular momentum vector, and + # axis 2 for coordinate + + norm = norm.reshape(1, -1) + norm = norm.T[:, :, np.newaxis] + return np.tensordot(prim_coeffs, norm * zeroth_part, (0, 0)).flatten() + +def _eval_nuc_deriv(coords, orders, center, angmom_comps, alphas, prim_coeffs, norm): + """ + Return the derivative of a Cartesian contraction with respect to nuclear coordinate. + + FixMe: this funciton works correctly only when M=1, L=1. + + Parameters + ---------- + coords : np.ndarray(N, 3) + Points in space where the derivative of the Gaussian primitive is evaluated. + Coordinates must be given as a two dimensional array. + orders : np.ndarray(3, ) + Orders of the derivative. + Has to be posive integer numbers greater than zero + 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(N, 3) + Evaluation of the derivative at each given coordinate. + """ + R_x, R_y, R_z = center + + # evaluating derivatives along each axis of the atomic coordinate + g_x = grad(eval_contractions, 1) + g_y = grad(eval_contractions, 2) + g_z = grad(eval_contractions, 3) + + # evaluating higher order derivatives + order_x, order_y, order_z = orders-1 # subtracting 1 because first order derivative has been calculated + for i in range(order_x): + g_x = grad(g_x, 1) + for j in range(order_y): + g_y = grad(g_y, 2) + for k in range(order_z): + g_z = grad(g_z, 3) + + # vectorizing derivatives to evaluate them at each coordinate efficiently + grad_x = np.apply_along_axis(g_x, 1, coords, R_x, R_y, R_z, angmom_comps, alphas, prim_coeffs, norm) + grad_y = np.apply_along_axis(g_y, 1, coords, R_x, R_y, R_z, angmom_comps, alphas, prim_coeffs, norm) + grad_z = np.apply_along_axis(g_z, 1, coords, R_x, R_y, R_z, angmom_comps, alphas, prim_coeffs, norm) + + return np.vstack([grad_x, grad_y, grad_z]).T + +def eval_nuc_deriv(coords, orders, center, angmom_comps, alphas, prim_coeffs, norms): + """ + Return the derivative of a Cartesian contraction with respect to nuclear coordinate. + + Parameters + ---------- + coords : np.ndarray(N, 3) + Points in space where the derivative of the Gaussian primitive is evaluated. + Coordinates must be given as a two dimensional array. + orders : np.ndarray(3, ) + Orders of the derivative. + Has to be posive integer numbers greater than zero + 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). + norms : np.ndarray(L, K) + Normalization constants for the primitives in each contraction. + + Returns + ------- + derivative : np.ndarray(M, L, N, 3) + Evaluation of the derivative at each given coordinate. + """ + M, L, N = prim_coeffs.shape[1], angmom_comps.shape[0], coords.shape[0] + output = np.zeros((M, L, N, 3)) + + for m in range(M): + prim_coeff = prim_coeffs[:, m] + + angular_output = np.zeros((L, N, 3)) + for l in range(L): + angmom_comp = angmom_comps[l] + norm = norms[l] + + derivative = _eval_nuc_deriv(coords, orders, center, angmom_comp, alphas, prim_coeff, norm) + angular_output[l] = derivative + output[m] = angular_output + return output \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 20f8aa96..2fb1a92e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,6 +35,7 @@ dependencies = [ "scipy>=1.4", "importlib_resources", "sympy", + "autograd" ] dynamic = ["version"] diff --git a/tests/test_ad.py b/tests/test_ad.py new file mode 100644 index 00000000..8caecc65 --- /dev/null +++ b/tests/test_ad.py @@ -0,0 +1,65 @@ +"""Test the ad module.""" +import pytest +import numpy as np +from gbasis.evals._deriv import _eval_deriv_contractions +from gbasis.gradient.ad import eval_nuc_deriv, _eval_nuc_deriv, eval_contractions + + +def test_eval_nuc_deriv(): + """Test the _eval_nuc_deriv function.""" + coords = np.random.rand(10, 3) + center = np.random.rand(3) + R_x, R_y, R_z = center + angmom_comps = np.random.randint(0, 2, (1, 3)) + alphas = np.random.rand(1) + + prim_coeffs = np.random.rand(2, 1) + norm = np.ones((1, 2)) + orders = np.array([1, 1, 1]) + + gradient = _eval_nuc_deriv(coords, orders, center, angmom_comps, alphas, prim_coeffs, norm) + + dx = 1e-5 + grad_numerical_x = (np.apply_along_axis(eval_contractions, 1, coords, R_x+dx, R_y, R_z, angmom_comps, alphas, prim_coeffs, norm) -\ + np.apply_along_axis(eval_contractions, 1, coords, R_x-dx, R_y, R_z, angmom_comps, alphas, prim_coeffs, norm))/(2*dx) + grad_numerical_y = (np.apply_along_axis(eval_contractions, 1, coords, R_x, R_y+dx, R_z, angmom_comps, alphas, prim_coeffs, norm) -\ + np.apply_along_axis(eval_contractions, 1, coords, R_x, R_y-dx, R_z, angmom_comps, alphas, prim_coeffs, norm))/(2*dx) + grad_numerical_z = (np.apply_along_axis(eval_contractions, 1, coords, R_x, R_y, R_z+dx, angmom_comps, alphas, prim_coeffs, norm) -\ + np.apply_along_axis(eval_contractions, 1, coords, R_x, R_y, R_z-dx, angmom_comps, alphas, prim_coeffs, norm))/(2*dx) + + grad_numerical = np.hstack([grad_numerical_x, grad_numerical_y, grad_numerical_z]) + + np.testing.assert_allclose(gradient, grad_numerical) + +def test_eval_nuc_deriv_full(): + """Test the eval_nuc_deriv function.""" + coords = np.random.rand(10, 3) + center = np.random.rand(3) + + R_x, R_y, R_z = center + angmom_comps = np.random.randint(0, 2, (2, 3)) + alphas = np.random.rand(4) + + prim_coeffs = np.random.rand(4, 5) + norms = np.ones((2, 4)) + + orders = np.array([1, 1, 1]) + output = eval_nuc_deriv(coords, orders, center, angmom_comps, alphas, prim_coeffs, norms) + + for i in range(3): + dx = np.zeros(3) + dx[i] = 1e-5 + grad_numerical = (_eval_deriv_contractions(coords, np.array([0, 0, 0]), center+dx, angmom_comps, alphas, prim_coeffs, norms) -\ + _eval_deriv_contractions(coords, np.array([0, 0, 0]), center-dx, angmom_comps, alphas, prim_coeffs, norms))/(2*dx[i]) + np.allclose(output[:, :, :, i], grad_numerical) + + # testing second order derivative + orders = np.array([2, 2, 2]) + output = eval_nuc_deriv(coords, orders, center, angmom_comps, alphas, prim_coeffs, norms) + for i in range(3): + dx = np.zeros(3) + dx[i] = 1e-5 + grad_numerical = (_eval_deriv_contractions(coords, np.array([0, 0, 0]), center+dx, angmom_comps, alphas, prim_coeffs, norms) +\ + _eval_deriv_contractions(coords, np.array([0, 0, 0]), center-dx, angmom_comps, alphas, prim_coeffs, norms) -\ + 2*_eval_deriv_contractions(coords, np.array([0, 0, 0]), center, angmom_comps, alphas, prim_coeffs, norms))/(dx[i]**2) + np.allclose(output[:, :, :, i], grad_numerical) \ No newline at end of file