Remove cdist_batched, which had a bug (#95)

superscreen/fem.py

@@ -5,7 +5,6 @@
 import scipy.linalg as la
 import scipy.sparse as sp
 from matplotlib.path import Path
-from scipy.spatial.distance import cdist
 
 
 def in_polygon(
@@ -417,28 +416,3 @@ def gradient_vertices(
         gx[i, :] = np.einsum("i, ij -> j", weights, Gx[t, :])
         gy[i, :] = np.einsum("i, ij -> j", weights, Gy[t, :])
     return sp.csr_matrix(gx), sp.csr_matrix(gy)
-
-
-def cdist_batched(
-    XA: np.ndarray, XB: np.ndarray, *, batch_size: int = 0, **kwargs
-) -> np.ndarray:
-    """A batched version of scipy.spatial.distance.cdist.
-
-    The batching is performed over the second input array, XB. See the docs
-    for scipy.spatial.distance.cdist for keyword argument options.
-
-    Args:
-        XA: An mA by n array of mA original observations in an n-dimensional space.
-        XB: An mB by n array of mB original observations in an n-dimensional space.
-        batch_size: The maximum size of each batch of XB values.
-            Setting batch_size <= 0 results in no batching.
-
-    Returns:
-        An mA by mB distance matrix.
-    """
-    if "out" in kwargs:
-        raise ValueError("cdist_batched does not support the 'out' keyword argument.")
-    if batch_size <= 0:
-        return cdist(XA, XB, **kwargs)
-    batches = np.array_split(XB, range(batch_size, XB.shape[0], batch_size))
-    return np.concatenate([cdist(XA, batch) for batch in batches], axis=1)

superscreen/solution.py

@@ -10,10 +10,11 @@
 import numpy as np
 import pint
 from scipy import interpolate
+from scipy.spatial.distance import cdist
 
 from .about import version_dict
 from .device import Device, Polygon
-from .fem import cdist_batched, in_polygon
+from .fem import in_polygon
 from .parameter import Constant
 from .sources.current import biot_savart_2d
 
@@ -737,7 +738,6 @@ def vector_potential_at_position(
         units: Optional[str] = None,
         with_units: bool = True,
         return_sum: bool = True,
-        batch_size: int = 0,
     ) -> Union[np.ndarray, Dict[str, np.ndarray]]:
         """Calculates the vector potential due to currents in the device at any
         point(s) in space. Note that this only considers the vector potential
@@ -767,8 +767,6 @@ def vector_potential_at_position(
                 with units attached.
             return_sum: Whether to return the sum of the potential from all layers in
                 the device, or a dict of ``{layer_name: potential_from_layer}``.
-            batch_size: The maximum size of each batch of positions.
-                See :func:`superscreen.fem.cdist_batched`.
 
         Returns:
             An np.ndarray if return_sum is True, otherwise a dict of
@@ -801,9 +799,7 @@ def vector_potential_at_position(
         if zs.ndim == 1:
             # We need zs to be shape (m, 1)
             zs = zs[:, np.newaxis]
-        rho2 = cdist_batched(
-            positions, points, batch_size=batch_size, metric="sqeuclidean"
-        ).astype(dtype, copy=False)
+        rho2 = cdist(positions, points, metric="sqeuclidean").astype(dtype, copy=False)
         # Compute the vector potential at the specified positions
         # from the currents in each layer
         vector_potentials = {}