Remove coord_type args from evals/integrals

gbasis/evals/density.py

@@ -62,7 +62,7 @@ def evaluate_density_using_evaluated_orbs(one_density_matrix, orb_eval):
     return np.sum(density, axis=0)
 
 
-def evaluate_density(one_density_matrix, basis, points, transform=None, coord_type="spherical"):
+def evaluate_density(one_density_matrix, basis, points, transform=None):
     r"""Return the density of the given basis set at the given points.
 
     Parameters
@@ -84,21 +84,14 @@ def evaluate_density(one_density_matrix, basis, points, transform=None, coord_ty
         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.
-    coord_type : {"cartesian", list/tuple of "cartesian" or "spherical", "spherical"}
-        Types of the coordinate system for the contractions.
-        If "cartesian", then all of the contractions are treated as Cartesian contractions.
-        If "spherical", then all of the contractions are treated as spherical contractions.
-        If list/tuple, then each entry must be a "cartesian" or "spherical" to specify the
-        coordinate type of each `GeneralizedContractionShell` instance.
-        Default value is "spherical".
 
     Returns
     -------
     density : np.ndarray(N,)
         Density evaluated at `N` grid points.
 
     """
-    orb_eval = evaluate_basis(basis, points, transform=transform, coord_type=coord_type)
+    orb_eval = evaluate_basis(basis, points, transform=transform)
     return evaluate_density_using_evaluated_orbs(one_density_matrix, orb_eval)
 
 
@@ -109,7 +102,6 @@ def evaluate_deriv_reduced_density_matrix(
     basis,
     points,
     transform=None,
-    coord_type="spherical",
 ):
     r"""Return the derivative of the first-order reduced density matrix at the given points.
 
@@ -156,13 +148,6 @@ def evaluate_deriv_reduced_density_matrix(
         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.
-    coord_type : {"cartesian", list/tuple of "cartesian" or "spherical", "spherical"}
-        Types of the coordinate system for the contractions.
-        If "cartesian", then all of the contractions are treated as Cartesian contractions.
-        If "spherical", then all of the contractions are treated as spherical contractions.
-        If list/tuple, then each entry must be a "cartesian" or "spherical" to specify the
-        coordinate type of each `GeneralizedContractionShell` instance.
-        Default value is "spherical".
 
     Returns
     -------
@@ -171,13 +156,13 @@ def evaluate_deriv_reduced_density_matrix(
 
     """
     deriv_orb_eval_one = evaluate_deriv_basis(
-        basis, points, orders_one, transform=transform, coord_type=coord_type
+        basis, points, orders_one, transform=transform
     )
     if np.array_equal(orders_one, orders_two):
         deriv_orb_eval_two = deriv_orb_eval_one
     else:
         deriv_orb_eval_two = evaluate_deriv_basis(
-            basis, points, orders_two, transform=transform, coord_type=coord_type
+            basis, points, orders_two, transform=transform
         )
     density = one_density_matrix.dot(deriv_orb_eval_two)
     density *= deriv_orb_eval_one
@@ -186,7 +171,7 @@ def evaluate_deriv_reduced_density_matrix(
 
 
 def evaluate_deriv_density(
-    orders, one_density_matrix, basis, points, transform=None, coord_type="spherical"
+    orders, one_density_matrix, basis, points, transform=None
 ):
     r"""Return the derivative of density of the given transformed basis set at the given points.
 
@@ -211,13 +196,6 @@ def evaluate_deriv_density(
         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.
-    coord_type : {"cartesian", list/tuple of "cartesian" or "spherical", "spherical"}
-        Types of the coordinate system for the contractions.
-        If "cartesian", then all of the contractions are treated as Cartesian contractions.
-        If "spherical", then all of the contractions are treated as spherical contractions.
-        If list/tuple, then each entry must be a "cartesian" or "spherical" to specify the
-        coordinate type of each `GeneralizedContractionShell` instance.
-        Default value is "spherical".
 
     Returns
     -------
@@ -250,14 +228,13 @@ def evaluate_deriv_density(
                     basis,
                     points,
                     transform=transform,
-                    coord_type=coord_type,
                 )
                 output += factor * num_occurence * density
     return output
 
 
 def evaluate_density_gradient(
-    one_density_matrix, basis, points, transform=None, coord_type="spherical"
+    one_density_matrix, basis, points, transform=None
 ):
     r"""Return the gradient of the density evaluated at the given points.
 
@@ -280,13 +257,6 @@ def evaluate_density_gradient(
         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.
-    coord_type : {"cartesian", list/tuple of "cartesian" or "spherical", "spherical"}
-        Types of the coordinate system for the contractions.
-        If "cartesian", then all of the contractions are treated as Cartesian contractions.
-        If "spherical", then all of the contractions are treated as spherical contractions.
-        If list/tuple, then each entry must be a "cartesian" or "spherical" to specify the
-        coordinate type of each `GeneralizedContractionShell` instance.
-        Default value is "spherical".
 
     Returns
     -------
@@ -302,30 +272,27 @@ def evaluate_density_gradient(
                 basis,
                 points,
                 transform=transform,
-                coord_type=coord_type,
             ),
             evaluate_deriv_density(
                 np.array([0, 1, 0]),
                 one_density_matrix,
                 basis,
                 points,
                 transform=transform,
-                coord_type=coord_type,
             ),
             evaluate_deriv_density(
                 np.array([0, 0, 1]),
                 one_density_matrix,
                 basis,
                 points,
                 transform=transform,
-                coord_type=coord_type,
             ),
         ]
     ).T
 
 
 def evaluate_density_laplacian(
-    one_density_matrix, basis, points, transform=None, coord_type="spherical"
+    one_density_matrix, basis, points, transform=None
 ):
     r"""Return the Laplacian of the density evaluated at the given points.
 
@@ -348,13 +315,6 @@ def evaluate_density_laplacian(
         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.
-    coord_type : {"cartesian", list/tuple of "cartesian" or "spherical", "spherical"}
-        Types of the coordinate system for the contractions.
-        If "cartesian", then all of the contractions are treated as Cartesian contractions.
-        If "spherical", then all of the contractions are treated as spherical contractions.
-        If list/tuple, then each entry must be a "cartesian" or "spherical" to specify the
-        coordinate type of each `GeneralizedContractionShell` instance.
-        Default value is "spherical".
 
     Returns
     -------
@@ -368,29 +328,26 @@ def evaluate_density_laplacian(
         basis,
         points,
         transform=transform,
-        coord_type=coord_type,
     )
     output += evaluate_deriv_density(
         np.array([0, 2, 0]),
         one_density_matrix,
         basis,
         points,
         transform=transform,
-        coord_type=coord_type,
     )
     output += evaluate_deriv_density(
         np.array([0, 0, 2]),
         one_density_matrix,
         basis,
         points,
         transform=transform,
-        coord_type=coord_type,
     )
     return output
 
 
 def evaluate_density_hessian(
-    one_density_matrix, basis, points, transform=None, coord_type="spherical"
+    one_density_matrix, basis, points, transform=None
 ):
     r"""Return the Hessian of the density evaluated at the given points.
 
@@ -413,13 +370,6 @@ def evaluate_density_hessian(
         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.
-    coord_type : {"cartesian", list/tuple of "cartesian" or "spherical", "spherical"}
-        Types of the coordinate system for the contractions.
-        If "cartesian", then all of the contractions are treated as Cartesian contractions.
-        If "spherical", then all of the contractions are treated as spherical contractions.
-        If list/tuple, then each entry must be a "cartesian" or "spherical" to specify the
-        coordinate type of each `GeneralizedContractionShell` instance.
-        Default value is "spherical".
 
     Returns
     -------
@@ -439,7 +389,6 @@ def evaluate_density_hessian(
                     basis,
                     points,
                     transform=transform,
-                    coord_type=coord_type,
                 )
                 for orders_one in np.identity(3, dtype=int)
             ]
@@ -449,7 +398,7 @@ def evaluate_density_hessian(
 
 
 def evaluate_posdef_kinetic_energy_density(
-    one_density_matrix, basis, points, transform=None, coord_type="spherical"
+    one_density_matrix, basis, points, transform=None
 ):
     r"""Return evaluations of positive definite kinetic energy density at the given points.
 
@@ -484,13 +433,6 @@ def evaluate_posdef_kinetic_energy_density(
         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.
-    coord_type : {"cartesian", list/tuple of "cartesian" or "spherical", "spherical"}
-        Types of the coordinate system for the contractions.
-        If "cartesian", then all of the contractions are treated as Cartesian contractions.
-        If "spherical", then all of the contractions are treated as spherical contractions.
-        If list/tuple, then each entry must be a "cartesian" or "spherical" to specify the
-        coordinate type of each `GeneralizedContractionShell` instance.
-        Default value is "spherical".
 
     Returns
     -------
@@ -508,14 +450,13 @@ def evaluate_posdef_kinetic_energy_density(
             basis,
             points,
             transform=transform,
-            coord_type=coord_type,
         )
     return 0.5 * output
 
 
 # TODO: test against a reference
 def evaluate_general_kinetic_energy_density(
-    one_density_matrix, basis, points, alpha, transform=None, coord_type="spherical"
+    one_density_matrix, basis, points, alpha, transform=None
 ):
     r"""Return evaluations of general form of the kinetic energy density at the given points.
 
@@ -546,13 +487,6 @@ def evaluate_general_kinetic_energy_density(
         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.
-    coord_type : {"cartesian", list/tuple of "cartesian" or "spherical", "spherical"}
-        Types of the coordinate system for the contractions.
-        If "cartesian", then all of the contractions are treated as Cartesian contractions.
-        If "spherical", then all of the contractions are treated as spherical contractions.
-        If list/tuple, then each entry must be a "cartesian" or "spherical" to specify the
-        coordinate type of each `GeneralizedContractionShell` instance.
-        Default value is "spherical".
 
     Returns
     -------
@@ -570,10 +504,10 @@ def evaluate_general_kinetic_energy_density(
         raise TypeError("`alpha` must be an int or float.")
 
     general_kinetic_energy_density = evaluate_posdef_kinetic_energy_density(
-        one_density_matrix, basis, points, transform=transform, coord_type=coord_type
+        one_density_matrix, basis, points, transform=transform
     )
     if alpha != 0:
         general_kinetic_energy_density += alpha * evaluate_density_laplacian(
-            one_density_matrix, basis, points, transform=transform, coord_type=coord_type
+            one_density_matrix, basis, points, transform=transform
         )
     return general_kinetic_energy_density

gbasis/evals/electrostatic_potential.py

@@ -10,7 +10,6 @@ def electrostatic_potential(
     nuclear_coords,
     nuclear_charges,
     transform=None,
-    coord_type="spherical",
     threshold_dist=0.0,
 ):
     r"""Return the electrostatic potentials of the basis set in the Cartesian form.
@@ -40,12 +39,6 @@ def electrostatic_potential(
         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.
-    coord_type : {"cartesian", "spherical", list/tuple of "cartesian" or "spherical"}
-        Types of the coordinate system for the contractions.
-        If "cartesian", then all of the contractions are treated as Cartesian contractions.
-        If "spherical", then all of the contractions are treated as spherical contractions.
-        If list/tuple, then each entry must be a "cartesian" or "spherical" to specify the
-        coordinate type of each `GeneralizedContractionShell` instance.
     threshold_dist : {float, 0.0}
         Threshold for rejecting nuclei whose distances to the points are less than the provided
         value. i.e. nuclei that are closer to the point than the threshold are discarded when
@@ -100,13 +93,15 @@ def electrostatic_potential(
     if threshold_dist < 0:
         raise ValueError("`threshold_dist` must be greater than or equal to zero.")
 
-    if coord_type == "cartesian":
+    coord_type = [type for type in [shell.coord_type for shell in basis]]
+
+    if all(type == "cartesian" for type in coord_type):
         if sum(cont.num_cart * cont.num_seg_cont for cont in basis) != one_density_matrix.shape[0]:
             raise ValueError(
                 "`one_density_matrix` does not have number of rows/columns that is equal to the "
                 "total number of Cartesian contractions (atomic orbitals)."
             )
-    elif coord_type == "spherical":
+    elif all(type == "spherical" for type in coord_type):
         if sum(cont.num_sph * cont.num_seg_cont for cont in basis) != one_density_matrix.shape[0]:
             raise ValueError(
                 "`one_density_matrix` does not have number of rows/columns that is equal to the "
@@ -128,10 +123,10 @@ def electrostatic_potential(
             )
     else:
         raise TypeError(
-            "`coord_type` must be 'spherical', 'cartesian', or a list/tuple of these strings."
+            "`coord_type` must be a list/tuple of the strings 'spherical' or 'cartesian'."
         )
     hartree_potential = point_charge_integral(
-        basis, points, -np.ones(points.shape[0]), transform=transform, coord_type=coord_type
+        basis, points, -np.ones(points.shape[0]), transform=transform
     )
     hartree_potential *= one_density_matrix[:, :, None]
     hartree_potential = np.sum(hartree_potential, axis=(0, 1))

gbasis/evals/eval.py

@@ -111,7 +111,7 @@ def construct_array_contraction(contractions, points):
         return output
 
 
-def evaluate_basis(basis, points, transform=None, coord_type="spherical"):
+def evaluate_basis(basis, points, transform=None):
     r"""Evaluate the basis set in the given coordinate system at the given points.
 
     Parameters
@@ -129,13 +129,6 @@ def evaluate_basis(basis, points, transform=None, coord_type="spherical"):
         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.
-    coord_type : {"cartesian", list/tuple of "cartesian" or "spherical", "spherical"}
-        Types of the coordinate system for the contractions.
-        If "cartesian", then all of the contractions are treated as Cartesian contractions.
-        If "spherical", then all of the contractions are treated as spherical contractions.
-        If list/tuple, then each entry must be a "cartesian" or "spherical" to specify the
-        coordinate type of each `GeneralizedContractionShell` instance.
-        Default value is "spherical".
 
     Returns
     -------
@@ -147,10 +140,12 @@ def evaluate_basis(basis, points, transform=None, coord_type="spherical"):
         `N` is the number of coordinates at which the contractions are evaluated.
 
     """
+    coord_type = [type for type in [shell.coord_type for shell in basis]]
+
     if transform is not None:
         return Eval(basis).construct_array_lincomb(transform, coord_type, points=points)
-    if coord_type == "cartesian":
+    if all(type == "cartesian" for type in coord_type):
         return Eval(basis).construct_array_cartesian(points=points)
-    if coord_type == "spherical":
+    if all(type == "spherical" for type in coord_type):
         return Eval(basis).construct_array_spherical(points=points)
     return Eval(basis).construct_array_mix(coord_type, points=points)