Add functions for 1rdms and 2-body exchange hole

gbasis/evals/density.py

@@ -95,12 +95,170 @@ def evaluate_density(one_density_matrix, basis, points, transform=None, threshol
 
  """
  orb_eval = evaluate_basis(basis, points, transform=transform)
- output = evaluate_density_using_evaluated_orbs(one_density_matrix, orb_eval)
+ density = evaluate_density_using_evaluated_orbs(one_density_matrix, orb_eval)
  # Fix #117: check magnitude of small negative density values, then use clip to remove them
- min_output = np.min(output)
- if min_output < 0.0 and abs(min_output) > threshold:
- raise ValueError(f"Found negative density <= {-threshold}, got {min_output}.")
- return output.clip(min=0.0)
+ min_density = np.min(density)
+ if min_density < 0.0 and abs(min_density) > threshold:
+ raise ValueError(f"Found negative density <= {-threshold}, got {min_density}.")
+ return density.clip(min=0.0)
+
+
+def evaluate_dm_using_evaluated_orbs(one_density_matrix, orb_eval_list):
+ """Return the evaluation of the density matrix given the evaluated orbitals.
+
+ Parameters
+ ----------
+ one_density_matrix : np.ndarray(K_orb, K_orb)
+ One-electron density matrix in terms of the given orbitals.
+ orb_eval_list : list of orbitals evaluated on grid points [np.ndarray(K_orb, N) ...]
+ Orbitals evaluated at :math:`N` grid points.
+ The set of orbitals must be the same basis set used to build the one-electron density
+ matrix.
+
+ Returns
+ -------
+ density : np.ndarray(N1,N2,K_orb,K_orb)
+ Density Matrix evaluated at `N1,N2` grid points.
+
+ Raises
+ ------
+ TypeError
+ If `orb_eval_list` does not consist of 2-dimensional `numpy` arrays with `dtype` float.
+ If `one_density_matrix` is not a symmetric, 2-dimensional `numpy` array with `dtype` float.
+
+ """
+ if not (
+ isinstance(one_density_matrix, np.ndarray)
+ and one_density_matrix.ndim == 2
+ and one_density_matrix.dtype == float
+ ):
+ raise TypeError(
+ "One-electron density matrix must be a two-dimensional `numpy` array with `dtype`"
+ " float."
+ )
+ for orb in orb_eval_list:
+ if not (isinstance(orb, np.ndarray) and orb.ndim == 2 and orb.dtype == float):
+ raise TypeError(
+ "Evaluation of orbitals must be a two-dimensional `numpy` array with `dtype` float."
+ )
+ if one_density_matrix.shape[0] != one_density_matrix.shape[1]:
+ raise ValueError("One-electron density matrix must be a square matrix.")
+
+ if not np.allclose(one_density_matrix, one_density_matrix.T):
+ raise ValueError("One-electron density matrix must be symmetric.")
+
+ # Tensor product for \gamma(\mathbf{r}_1,\mathbf{r}_2) = \sum_{pq} \gamma_{pq} \chi_p(\mathbf{r}_1) \chi_q(\mathbf{r}_2)
+ density = np.einsum("ij,ik,jl->klij", one_density_matrix, orb_eval_list[0], orb_eval_list[1])
+
+ # returns dm evaluated on each grid point
+ return density
+
+
+def evaluate_dm_density(one_density_matrix, basis, points_list, transform=None):
+ r"""Return the density of the given basis set at the given points.
+
+ Parameters
+ ----------
+ one_density_matrix : np.ndarray(K_orbs, K_orbs)
+ One-electron density matrix in terms of the given basis set.
+ If the basis is transformed using `transform` keyword, then the density matrix is assumed to
+ be expressed with respect to the transformed basis set.
+ basis : list/tuple of GeneralizedContractionShell
+ Shells of generalized contractions.
+ points_list : list of points [np.ndarray(N, 3)...]
+ Cartesian coordinates of the points in space (in atomic units) where the basis functions
+ are evaluated.
+ Rows correspond to the points and columns correspond to the :math:`x, y, \text{and} z`
+ components.
+ This function can take a list of points at which basis functions are evaluated. If only one
+ set of points is given, it will be duplicated.
+ transform : np.ndarray(K_orbs, K_cont)
+ Transformation matrix from the basis set in the given coordinate system (e.g. AO) to linear
+ combinations of contractions (e.g. MO).
+ Transformation is applied to the left, i.e. the sum is over the index 1 of `transform`
+ and index 0 of the array for contractions.
+ Default is no transformation.
+
+
+ Returns
+ -------
+ dm_on_grid : np.ndarray(N1,N2,K_orb,K_orb)
+ Density Matrix evaluated at `N1,N2` grid points.
+
+ """
+
+ orb_evals = []
+ # evaluate basi(e)s on the point set(s)
+ for grid in points_list:
+ orb_eval = evaluate_basis(basis, grid, transform=transform)
+ orb_evals.append(orb_eval)
+ # if only one set of points is supplied, it is duplicated
+ if len(points_list) == 1:
+ orb_evals.append(orb_eval)
+
+ # Calulated performed using the evaluated orbitals
+ dm_on_grid = evaluate_dm_using_evaluated_orbs(one_density_matrix, orb_evals)
+
+ return dm_on_grid
+
+
+def evaluate_hole_x2(one_density_matrix, basis, points_list, transform=None):
+ r"""Return the two-body hole correlation function.
+
+ .. math ::
+
+ \left.
+ h(\mathbf{r}_1,\mathbf{r}_2) =
+ \right.
+ -\frac{|\gamma(\mathbf{r}_1,\mathbf{r}_2)|^2}
+ {\rho({\mathbf{r}_1})\rho({\mathbf{r}_2})}
+
+ Parameters
+ ----------
+ one_density_matrix : np.ndarray(K_orbs, K_orbs)
+ One-electron density matrix in terms of the given basis set.
+ If the basis is transformed using `transform` keyword, then the density matrix is assumed to
+ be expressed with respect to the transformed basis set.
+ basis : list/tuple of GeneralizedContractionShell
+ Shells of generalized contractions.
+ points_list : list of points [np.ndarray(N, 3)]
+ Cartesian coordinates of the points in space (in atomic units) where the basis functions
+ are evaluated.
+ Rows correspond to the points and columns correspond to the :math:`x, y, \text{and} z`
+ components.
+ This function can take a list of points at which basis functions are evaluated. If only one
+ set of points is given, it will be duplicated.
+ transform : np.ndarray(K_orbs, K_cont)
+ Transformation matrix from the basis set in the given coordinate system (e.g. AO) to linear
+ combinations of contractions (e.g. MO).
+ Transformation is applied to the left, i.e. the sum is over the index 1 of `transform`
+ and index 0 of the array for contractions.
+ Default is no transformation.
+
+
+ Returns
+ -------
+ hole_x2 : np.ndarray(N1,N2,K_orb,K_orb)
+ Two-body Exchange Hole evaluated at `N1,N2` grid points.
+
+ """
+ dens_evals = []
+
+ for grid in points_list:
+ dens_eval = evaluate_density(one_density_matrix, basis, grid, transform=transform)
+ dens_evals.append(dens_eval)
+
+ if len(points_list) == 1:
+ dens_evals.append(dens_eval)
+
+ # build density matrix on grid
+ dm_eval = evaluate_dm_density(one_density_matrix, basis, points_list, transform=transform)
+
+ # evaluate hole function
+ numerator = np.einsum("ijkl,ijkl->ijkl", dm_eval, dm_eval)
+ hole_x2 = -1 * np.einsum("ijkl,i,j->ijkl", numerator, 1 / dens_evals[0], 1 / dens_evals[1])
+
+ return hole_x2
 
 
 def evaluate_deriv_reduced_density_matrix(
@@ -188,8 +346,8 @@ def evaluate_deriv_reduced_density_matrix(
  )
  density = one_density_matrix.dot(deriv_orb_eval_two)
  density *= deriv_orb_eval_one
- density = np.sum(density, axis=0)
- return density
+ deriv_reduced_density_matrix = np.sum(density, axis=0)
+ return deriv_reduced_density_matrix
 
 
 def evaluate_deriv_density(
@@ -238,7 +396,7 @@ def evaluate_deriv_density(
  # pylint: disable=R0914
  total_l_x, total_l_y, total_l_z = orders
 
- output = np.zeros(points.shape[0])
+ density_deriv = np.zeros(points.shape[0])
  for l_x in range(total_l_x // 2 + 1):
  # prevent double counting for the middle of the even total_l_x
  # e.g. If total_l_x == 4, then l_x is in [0, 1, 2, 3, 4]. Exploiting symmetry we only need
@@ -273,8 +431,8 @@ def evaluate_deriv_density(
  transform=transform,
  deriv_type=deriv_type,
  )
- output += factor * num_occurence * density
- return output
+ density_deriv += factor * num_occurence * density
+ return density_deriv
 
 
 def evaluate_density_gradient(
@@ -318,7 +476,7 @@ def evaluate_density_gradient(
 
  """
  orders_one = np.array(([1, 0, 0], [0, 1, 0], [0, 0, 1]))
- output = np.zeros((3, len(points)))
+ density_gradient = np.zeros((3, len(points)))
  # Evaluation of generalized contraction shell for zeroth order = 0,0,0
  zeroth_deriv = evaluate_deriv_basis(
  basis,
@@ -340,8 +498,8 @@ def evaluate_density_gradient(
  # output[ind] = 2*(np.einsum('ij,ik,jk -> k',one_density_matrix, zeroth_deriv, deriv_comp))
  density = one_density_matrix.dot(zeroth_deriv)
  density *= deriv_comp
- output[ind] = 2 * 1 * np.sum(density, axis=0)
- return output.T
+ density_gradient[ind] = 2 * 1 * np.sum(density, axis=0)
+ return density_gradient.T
 
 
 def evaluate_density_laplacian(
@@ -387,7 +545,7 @@ def evaluate_density_laplacian(
  orders_one_second = np.array(([2, 0, 0], [0, 2, 0], [0, 0, 2]))
  orders_one_first = np.array(([1, 0, 0], [0, 1, 0], [0, 0, 1]))
  orders_two = np.array(([1, 0, 0], [0, 1, 0], [0, 0, 1]))
- output = np.zeros(points.shape[0])
+ density_laplacian = np.zeros(points.shape[0])
  # Evaluation of generalized contraction shell for zeroth order = 0,0,0
  zeroth_deriv = evaluate_deriv_basis(
  basis,
@@ -409,7 +567,7 @@ def evaluate_density_laplacian(
 
  density = one_density_matrix.dot(zeroth_deriv)
  density *= deriv_one
- output += 2 * 1 * np.sum(density, axis=0)
+ density_laplacian += 2 * 1 * np.sum(density, axis=0)
 
  for orders in zip(orders_one_first, orders_two):
  deriv_one = evaluate_deriv_basis(
@@ -430,9 +588,9 @@ def evaluate_density_laplacian(
  # output[ind] = 2*(np.einsum('ij,ik,jk -> k',one_density_matrix, zeroth_deriv, deriv_comp))
  density = one_density_matrix.dot(deriv_two)
  density *= deriv_one
- output += 2 * 1 * np.sum(density, axis=0)
+ density_laplacian += 2 * 1 * np.sum(density, axis=0)
 
- return output
+ return density_laplacian
 
 
 def evaluate_density_hessian(
@@ -545,13 +703,13 @@ def evaluate_density_hessian(
  density_2 = np.einsum("ijkm,ijmk -> ijkm", one_two_arr_1, raw_density_2)
 
  # factors and sum over basis functions
- output = 2 * 1 * np.sum(density_1, axis=2)
- output += 2 * 1 * np.sum(density_2, axis=2)
+ density_hessian = 2 * 1 * np.sum(density_1, axis=2)
+ density_hessian += 2 * 1 * np.sum(density_2, axis=2)
 
  # copying lower matrix to upper matrix
- upp = np.swapaxes(output, 0, 1)
+ upp = np.swapaxes(density_hessian, 0, 1)
  upp = np.triu(upp.T, 1)
- return output.T + upp
+ return density_hessian.T + upp
 
 
 def evaluate_posdef_kinetic_energy_density(
@@ -610,9 +768,9 @@ def evaluate_posdef_kinetic_energy_density(
  `N` grid points.
 
  """
- output = np.zeros(points.shape[0])
+ posdef_kindetic_energy_density = np.zeros(points.shape[0])
  for orders in np.identity(3, dtype=int):
- output += evaluate_deriv_reduced_density_matrix(
+ posdef_kindetic_energy_density += evaluate_deriv_reduced_density_matrix(
  orders,
  orders,
  one_density_matrix,
@@ -625,7 +783,7 @@ def evaluate_posdef_kinetic_energy_density(
  min_output = np.min(output)
  if min_output < 0.0 and abs(min_output) > threshold:
  raise ValueError(f"Found negative density <= {-threshold}, got {min_output}.")
- return (0.5 * output).clip(min=0.0)
+ return (0.5 * posdef_kindetic_energy_density).clip(min=0.0)
 
 
 # TODO: test against a reference

tests/test_density.py

@@ -1,6 +1,9 @@
 """Test gbasis.evals.density."""
 from gbasis.evals.density import (
  evaluate_density,
+ evaluate_dm_using_evaluated_orbs,
+ evaluate_dm_density,
+ evaluate_hole_x2,
  evaluate_density_gradient,
  evaluate_density_hessian,
  evaluate_density_laplacian,
@@ -74,6 +77,57 @@ def test_evaluate_density():
  assert np.allclose(dens, np.einsum("ij,ik,jk->k", density, evaluate_orbs, evaluate_orbs))
 
 
+def test_evaluate_dm_using_evaluated_orbs():
+ """Test gbasis.evals.density.evaluate_density_using_evaluated_orbs."""
+ density_mat = np.array([[1.0, 2.0], [2.0, 3.0]])
+ orb_eval = np.array([[1.0], [2.0]])
+ dens = evaluate_dm_using_evaluated_orbs(density_mat, [orb_eval, orb_eval])
+ assert np.all(dens >= 0.0)
+ assert np.allclose(
+ np.einsum("iikl->i", dens), np.einsum("ij,ik,jk->k", density_mat, orb_eval, orb_eval)
+ )
+
+ density_mat = np.array([[1.0, 2.0], [2.0, 3.0]])
+ orb_eval = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
+ dens = evaluate_dm_using_evaluated_orbs(density_mat, [orb_eval, orb_eval])
+ assert np.all(dens >= 0)
+ assert np.allclose(
+ np.einsum("iikl->i", dens), np.einsum("ij,ik,jk->k", density_mat, orb_eval, orb_eval)
+ )
+
+
+def test_evaluate_dm_density():
+ """Test gbasis.evals.density.evaluate_density."""
+ basis_dict = parse_nwchem(find_datafile("data_sto6g.nwchem"))
+ basis = make_contractions(basis_dict, ["Kr"], np.array([[0, 0, 0]]), "spherical")
+ transform = np.random.rand(14, 18)
+ density = np.random.rand(14, 14)
+ density += density.T
+ points = np.random.rand(100, 3)
+
+ evaluate_orbs = evaluate_basis(basis, points, transform)
+ dens = evaluate_dm_density(density, basis, [points], transform)
+ assert np.allclose(
+ np.einsum("iikl->i", dens), np.einsum("ij,ik,jk->k", density, evaluate_orbs, evaluate_orbs)
+ )
+
+
+def test_evaluate_hole_x2():
+ basis_dict = parse_nwchem(find_datafile("data_sto6g.nwchem"))
+ basis = make_contractions(basis_dict, ["Kr"], np.array([[0, 0, 0]]), "spherical")
+ transform = np.ones((18, 18))
+ density = np.load(find_datafile("data_rdm_kr_sto6g.npy"))
+
+ points = np.random.rand(10, 3)
+
+ eh = evaluate_hole_x2(density, basis, points_list=[points, points], transform=transform)
+ ed = evaluate_density(density, basis, points, transform=transform)
+
+ assert np.allclose(
+ np.einsum("ijkl,j->", eh, ed) / np.sum(ed) / len(points) * len(transform), -1
+ )
+
+
 def test_evaluate_deriv_density():
  """Test gbasis.evals.density.evaluate_deriv_density."""
  basis_dict = parse_nwchem(find_datafile("data_sto6g.nwchem"))
@@ -448,3 +502,6 @@ def test_evaluate_general_kinetic_energy_density():
  evaluate_posdef_kinetic_energy_density(np.identity(40), basis, points, np.identity(40))
  + evaluate_density_laplacian(np.identity(40), basis, points, np.identity(40)),
  )
+
+
+test_evaluate_hole_x2()