nicer implementation

nifty/spaces/lm_space/lm_space.py

@@ -169,6 +169,12 @@ class LMSpace(Space):
 
         return dists
 
+    def get_unique_distances(self):
+        return np.arange(self.lmax+1, dtype=np.float64)
+
+    def get_natural_binbounds(self):
+        return np.arange(self.lmax, dtype=np.float64) + 0.5
+
     @staticmethod
     def _distance_array_helper(index_array, lmax):
         u = 2*lmax + 1

nifty/spaces/power_space/power_space.py

@@ -95,9 +95,6 @@ class PowerSpace(Space):
             raise ValueError("harmonic_partner must be a harmonic space.")
         self._harmonic_partner = harmonic_partner
 
-        dists = self._harmonic_partner.get_distance_array(
-                distribution_strategy)
-
         # sanity check
         if binbounds is not None and not(nbin is None and logarithmic is None):
             raise ValueError(
@@ -105,16 +102,7 @@ class PowerSpace(Space):
 
         self._binbounds = None
         if logarithmic is None and nbin is None and binbounds is None:
-            # compute the "natural" binning, i.e. there
-            # is one bin for every distinct k-vector length
-            tmp = dists.unique()  # expensive!
-            tol = 1e-12*tmp[-1]
-            # remove all points that are closer than tol to their right
-            # neighbors.
-            # I'm appending the last value*2 to the array to treat the
-            # rightmost point correctly.
-            tmp = tmp[np.diff(np.r_[tmp, 2*tmp[-1]]) > tol]
-            bb = tmp[0:-1]+0.5*np.diff(tmp)
+            bb = self._harmonic_partner.get_natural_binbounds()
         else:
             if binbounds is not None:
                 bb = np.sort(np.array(binbounds))
@@ -126,7 +114,7 @@ class PowerSpace(Space):
 
                 # equidistant binning (linear or log)
                 # MR FIXME: this needs to improve
-                kindex = dists.unique()
+                kindex = self._harmonic_partner.get_unique_distances()
                 if (logarithmic):
                     k = np.r_[0, np.log(kindex[1:])]
                 else:
@@ -145,6 +133,8 @@ class PowerSpace(Space):
                     bb = np.exp(bb)
             self._binbounds = tuple(bb)
 
+        dists = self._harmonic_partner.get_distance_array(
+                distribution_strategy)
         # Compute pindex, kindex and rho according to bb
         self._pindex = distributed_data_object(
             global_shape=dists.shape,

nifty/spaces/rg_space/rg_space.py

@@ -283,6 +283,36 @@ class RGSpace(Space):
         dists = np.sqrt(dists)
         return dists
 
+    def get_unique_distances(self):
+        dimensions = len(self.shape)
+        if dimensions == 1:  # extra easy
+            maxdist = self.shape[0]//2
+            return np.arange(maxdist+1, dtype=np.float64) * self.distances[0]
+        if np.all(self.distances == self.distances[0]):  # shortcut
+            maxdist = np.asarray(self.shape)//2
+            tmp = np.sum(maxdist*maxdist)
+            tmp = np.zeros(tmp+1, dtype=np.bool)
+            t2 = np.arange(maxdist[0]+1, dtype=np.int64)
+            t2 *= t2
+            for i in range(1, dimensions):
+                t3 = np.arange(maxdist[i]+1, dtype=np.int64)
+                t3 *= t3
+                t2 = np.add.outer(t2, t3)
+            tmp[t2] = True
+            return np.sqrt(np.nonzero(tmp)[0])*self.distances[0]
+        else:  # do it the hard way
+            tmp = self.get_distance_array('not').unique()  # expensive!
+            tol = 1e-12*tmp[-1]
+            # remove all points that are closer than tol to their right
+            # neighbors.
+            # I'm appending the last value*2 to the array to treat the
+            # rightmost point correctly.
+            return tmp[np.diff(np.r_[tmp, 2*tmp[-1]]) > tol]
+
+    def get_natural_binbounds(self):
+        tmp = self.get_unique_distances()
+        return 0.5*(tmp[:-1]+tmp[1:])
+
     def get_fft_smoothing_kernel_function(self, sigma):
         return lambda x: np.exp(-0.5 * np.pi*np.pi * x*x * sigma*sigma)
 

nifty/spaces/space/space.py

@@ -117,6 +117,20 @@ class Space(DomainObject):
         raise NotImplementedError(
             "There is no generic distance structure for Space base class.")
 
+    def get_unique_distances(self):
+        raise NotImplementedError
+
+    def get_natural_binbounds(self):
+        """ The boundaries for natural power spectrum binning.
+
+        Returns
+        -------
+        distributed_data_object
+            A numpy array containing the binbounds
+
+        """
+        raise NotImplementedError
+
     def get_fft_smoothing_kernel_function(self, sigma):
         """ This method returns a smoothing kernel function.