From 2afa15a72e4221802dc29c9de9bd9b90f77f9a7a Mon Sep 17 00:00:00 2001
From: Logan Bishop-Van Horn <logan.bvh@gmail.com>
Date: Sat, 6 May 2023 10:43:51 -0700
Subject: [PATCH] Remove cdist_batched, which had a bug (#95)

---
 superscreen/fem.py      | 26 --------------------------
 superscreen/solution.py | 10 +++-------
 2 files changed, 3 insertions(+), 33 deletions(-)

diff --git a/superscreen/fem.py b/superscreen/fem.py
index 428e44ab..6e88eddb 100644
--- a/superscreen/fem.py
+++ b/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)
diff --git a/superscreen/solution.py b/superscreen/solution.py
index e699857e..3a9d23f8 100644
--- a/superscreen/solution.py
+++ b/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 = {}