Merge pull request #317 from python-adaptive/to_numpy

add to_numpy methods

Author: jbweston

adaptive/learner/average_learner.py

@@ -56,6 +56,10 @@ def __init__(self, function, atol=None, rtol=None, min_npoints=2):
     def n_requested(self):
         return self.npoints + len(self.pending_points)
 
+    def to_numpy(self):
+        """Data as NumPy array of size (npoints, 2) with seeds and values."""
+        return np.array(sorted(self.data.items()))
+
     def ask(self, n, tell_pending=True):
         points = list(range(self.n_requested, self.n_requested + n))
 

adaptive/learner/integrator_learner.py

@@ -546,6 +546,10 @@ def plot(self):
         xs, ys = zip(*[(x, y) for ival in ivals for x, y in sorted(ival.data.items())])
         return hv.Path((xs, ys))
 
+    def to_numpy(self):
+        """Data as NumPy array of size (npoints, 2)."""
+        return np.array(sorted(self.data.items()))
+
     def _get_data(self):
         # Change the defaultdict of SortedSets to a normal dict of sets.
         x_mapping = {k: set(v) for k, v in self.x_mapping.items()}

adaptive/learner/learner1D.py

@@ -269,6 +269,11 @@ def vdim(self):
                 return 1
         return self._vdim
 
+    def to_numpy(self):
+        """Data as NumPy array of size ``(npoints, 2)`` if ``learner.function`` returns a scalar
+        and ``(npoints, 1+vdim)`` if ``learner.function`` returns a vector of length ``vdim``."""
+        return np.array([(x, *np.atleast_1d(y)) for x, y in sorted(self.data.items())])
+
     @property
     def npoints(self):
         """Number of evaluated points."""

adaptive/learner/learner2D.py

@@ -377,6 +377,13 @@ def xy_scale(self):
         else:
             return np.array([xy_scale[0], xy_scale[1] / self.aspect_ratio])
 
+    def to_numpy(self):
+        """Data as NumPy array of size ``(npoints, 3)`` if ``learner.function`` returns a scalar
+        and ``(npoints, 2+vdim)`` if ``learner.function`` returns a vector of length ``vdim``."""
+        return np.array(
+            [(x, y, *np.atleast_1d(z)) for (x, y), z in sorted(self.data.items())]
+        )
+
     def _scale(self, points):
         points = np.asarray(points, dtype=float)
         return (points - self.xy_mean) / self.xy_scale

adaptive/learner/learnerND.py

@@ -379,6 +379,12 @@ def vdim(self):
                 self._vdim = 1
         return self._vdim if self._vdim is not None else 1
 
+    def to_numpy(self):
+        """Data as NumPy array of size ``(npoints, dim+vdim)``, where ``dim`` is the
+        size of the input dimension and ``vdim`` is the length of the return value
+        of ``learner.function``."""
+        return np.array([(*p, *np.atleast_1d(v)) for p, v in sorted(self.data.items())])
+
     @property
     def bounds_are_done(self):
         return all(p in self.data for p in self._bounds_points)