many tweaks

nifty/data_objects/distributed_do.py

@@ -22,8 +22,6 @@ def get_locshape(shape, distaxis):
         distaxis = -1
     if distaxis==-1:
         return shape
-    if distaxis<0 or distaxis>=len(shape):
-        print distaxis,shape
     shape2=list(shape)
     shape2[distaxis]=shareSize(shape[distaxis],ntask,rank)
     return tuple(shape2)
@@ -113,9 +111,7 @@ class data_object(object):
             for ax in axis:
                 if ax<self._distaxis:
                     shift+=1
-            print (axis,self._distaxis,shift)
             shp[self._distaxis-shift] = self.shape[self._distaxis]
-            print (self.shape, shp)
             return from_local_data(shp, res, self._distaxis-shift)
 
         # check if the result is scalar or if a result_field must be constr.
@@ -134,7 +130,6 @@ class data_object(object):
             if a._shape != b._shape:
                 raise ValueError("shapes are incompatible.")
             if a._distaxis != b._distaxis:
-                print (a._distaxis, b._distaxis)
                 raise ValueError("distributions are incompatible.")
             a = a._data
             b = b._data
@@ -286,14 +281,6 @@ def from_random(random_type, shape, dtype=np.float64, distaxis=0, **kwargs):
     return data_object(shape, generator_function(dtype=dtype, shape=lshape, **kwargs), distaxis=distaxis)
 
 
-def to_ndarray(arr):
-    return arr._data
-
-
-def from_ndarray(arr,distaxis=0):
-    return data_object(arr.shape,arr,distaxis)
-
-
 def local_data(arr):
     return arr._data
 
@@ -350,6 +337,7 @@ def redistribute (arr, dist=None, nodist=None):
             if i not in nodist:
                 dist=i
                 break
+
     if arr._distaxis==-1:  # just pick the proper subset
         return from_global_data(arr._data, dist)
     if dist==-1: # gather data
@@ -363,12 +351,17 @@ def redistribute (arr, dist=None, nodist=None):
         disp[1:]=np.cumsum(sz[:-1])
         tmp=tmp.flatten()
         out = np.empty(arr.size,dtype=arr.dtype)
-        print tmp.shape, out.shape, sz, disp
         comm.Allgatherv(tmp,[out,sz,disp,MPI.BYTE])
-        out = out.reshape(arr._shape)
+        shp = np.array(arr._shape)
+        shp[1:arr._distaxis+1] = shp[0:arr._distaxis]
+        shp[0] = arr.shape[arr._distaxis]
+        out = out.reshape(shp)
         out = np.moveaxis(out, 0, arr._distaxis)
         return from_global_data (out, distaxis=-1)
     # real redistribution via Alltoallv
+    # temporary slow, but simple solution
+    return redistribute(redistribute(arr,dist=-1),dist=dist)
+
     tmp = np.moveaxis(arr._data, (dist, arr._distaxis), (0, 1))
     tshape = tmp.shape
     slabsize=np.prod(tmp.shape[2:])*tmp.itemsize
@@ -383,7 +376,6 @@ def redistribute (arr, dist=None, nodist=None):
     rdisp[0]=0
     sdisp[1:]=np.cumsum(ssz[:-1])
     rdisp[1:]=np.cumsum(rsz[:-1])
-    print ssz, rsz
     tmp=tmp.flatten()
     out = np.empty(np.prod(get_locshape(arr.shape,dist)),dtype=arr.dtype)
     s_msg = [tmp, (ssz, sdisp), MPI.BYTE]

nifty/data_objects/my_own_do.py

@@ -208,14 +208,6 @@ def from_random(random_type, shape, dtype=np.float64, **kwargs):
     return data_object(generator_function(dtype=dtype, shape=shape, **kwargs))
 
 
-def to_ndarray(arr):
-    return arr._data
-
-
-def from_ndarray(arr):
-    return data_object(arr)
-
-
 def local_data(arr):
     return arr._data
 
@@ -233,16 +225,12 @@ def distaxis(arr):
 
 
 def from_local_data (shape, arr, distaxis):
-    if distaxis!=-1:
-        raise NotImplementedError
     if shape!=arr.shape:
         raise ValueError
     return data_object(arr)
 
 
 def from_global_data (arr, distaxis=-1):
-    if distaxis!=-1:
-        raise NotImplementedError
     return data_object(arr)
 
 
@@ -251,8 +239,6 @@ def to_global_data (arr):
 
 
 def redistribute (arr, dist=None, nodist=None):
-    if dist is not None and dist!=-1:
-        raise NotImplementedError
     return arr
 
 
@@ -261,6 +247,4 @@ def default_distaxis():
 
 
 def local_shape(glob_shape, distaxis):
-    if distaxis!=-1:
-        raise NotImplementedError
     return glob_shape

nifty/data_objects/numpy_do.py

@@ -15,14 +15,6 @@ def from_random(random_type, shape, dtype=np.float64, **kwargs):
     return generator_function(dtype=dtype, shape=shape, **kwargs)
 
 
-def to_ndarray(arr):
-    return arr
-
-
-def from_ndarray(arr):
-    return np.asarray(arr)
-
-
 def local_data(arr):
     return arr
 
@@ -40,16 +32,12 @@ def distaxis(arr):
 
 
 def from_local_data (shape, arr, distaxis):
-    if distaxis!=-1:
-        raise NotImplementedError
     if shape!=arr.shape:
         raise ValueError
     return arr
 
 
 def from_global_data (arr, distaxis=-1):
-    if distaxis!=-1:
-        raise NotImplementedError
     return arr
 
 
@@ -58,8 +46,6 @@ def to_global_data (arr):
 
 
 def redistribute (arr, dist=None, nodist=None):
-    if dist is not None and dist!=-1:
-        raise NotImplementedError
     return arr
 
 
@@ -68,6 +54,4 @@ def default_distaxis():
 
 
 def local_shape(glob_shape, distaxis):
-    if distaxis!=-1:
-        raise NotImplementedError
     return glob_shape

nifty/field.py

@@ -313,6 +313,9 @@ class Field(object):
                 new_shape[self.domain.axes[ind][0]:
                           self.domain.axes[ind][-1]+1] = wgt.shape
                 wgt = wgt.reshape(new_shape)
+                # FIXME only temporary
+                if ind==0:
+                    wgt = dobj.local_data(dobj.from_global_data(wgt, distaxis=0))
                 out *= wgt**power
         fct = fct**power
         if fct != 1.:

nifty/operators/diagonal_operator.py

@@ -22,7 +22,7 @@ from ..field import Field
 from ..domain_tuple import DomainTuple
 from .endomorphic_operator import EndomorphicOperator
 from ..nifty_utilities import cast_iseq_to_tuple
-from ..dobj import to_ndarray as to_np
+from .. import dobj
 
 
 class DiagonalOperator(EndomorphicOperator):
@@ -152,5 +152,7 @@ class DiagonalOperator(EndomorphicOperator):
         if self._spaces is None:
             return diag*x
 
-        reshaped_local_diagonal = np.reshape(to_np(diag.val), self._reshaper)
+        reshaped_local_diagonal = np.reshape(dobj.to_global_data(diag.val), self._reshaper)
+        if 0 in self._spaces:
+            reshaped_local_diagonal = dobj.local_data(dobj.from_global_data(reshaped_local_diagonal,distaxis=0))
         return Field(x.domain, val=x.val*reshaped_local_diagonal)

nifty/operators/fft_operator_support.py

@@ -138,12 +138,13 @@ class SphericalTransformation(Transformation):
                 odat[slice] = self._slice_p2h(idat[slice])
             odat = dobj.from_local_data(self.hdom.shape, odat, distaxis)
             if distaxis!= dobj.distaxis(x.val):
-                odat = dobj.redistribute(odat, dist=distaxis)
+                odat = dobj.redistribute(odat, dist=dobj.distaxis(x.val))
             return Field(self.hdom, odat)
         else:
-            res = Field(self.pdom, dtype=x.dtype)
-            odat = dobj.local_data(res.val)
-
+            odat = np.empty(dobj.local_shape(self.pdom.shape, distaxis=distaxis), dtype=x.dtype)
             for slice in utilities.get_slice_list(idat.shape, axes):
                 odat[slice] = self._slice_h2p(idat[slice])
-        return res
+            odat = dobj.from_local_data(self.pdom.shape, odat, distaxis)
+            if distaxis!= dobj.distaxis(x.val):
+                odat = dobj.redistribute(odat, dist=dobj.distaxis(x.val))
+            return Field(self.pdom, odat)

nifty/operators/laplace_operator.py

@@ -21,7 +21,7 @@ from ..field import Field
 from ..spaces.power_space import PowerSpace
 from .endomorphic_operator import EndomorphicOperator
 from .. import DomainTuple
-from ..dobj import to_ndarray as to_np, from_ndarray as from_np
+from .. import dobj
 
 
 class LaplaceOperator(EndomorphicOperator):
@@ -89,9 +89,13 @@ class LaplaceOperator(EndomorphicOperator):
         return self._logarithmic
 
     def _times(self, x):
-        val = to_np(x.val)
+        val = dobj.to_global_data(x.val)
         axes = x.domain.axes[self._space]
         axis = axes[0]
+        locval = x.val
+        if axis == dobj.distaxis(locval):
+            locval = dobj.redistribute(locval, nodist=(axis,))
+        val = dobj.local_data(locval)
         nval = len(self._dposc)
         prefix = (slice(None),) * axis
         sl_l = prefix + (slice(None, -1),)  # "left" slice
@@ -106,8 +110,10 @@ class LaplaceOperator(EndomorphicOperator):
         ret /= np.sqrt(dposc)
         ret[prefix + (slice(None, 2),)] = 0.
         ret[prefix + (-1,)] = 0.
-        return Field(self.domain, val=from_np(ret)).weight(-0.5,
-                                                           spaces=self._space)
+        ret = dobj.from_local_data(locval.shape, ret, dobj.distaxis(locval))
+        if dobj.distaxis(locval)!=dobj.distaxis(x.val):
+            ret = dobj.redistribute(ret, dist=dobj.distaxis(x.val))
+        return Field(self.domain, val=ret).weight(-0.5, spaces=self._space)
 
     def _adjoint_times(self, x):
         axes = x.domain.axes[self._space]
@@ -118,7 +124,10 @@ class LaplaceOperator(EndomorphicOperator):
         sl_r = prefix + (slice(1, None),)  # "right" slice
         dpos = self._dpos.reshape((1,)*axis + (nval-1,))
         dposc = self._dposc.reshape((1,)*axis + (nval,))
-        y = to_np(x.weight(0.5, spaces=self._space).val)
+        yf = x.weight(0.5, spaces=self._space).val
+        if axis == dobj.distaxis(yf):
+            yf = dobj.redistribute(yf, nodist=(axis,))
+        y=dobj.local_data(yf)
         y /= np.sqrt(dposc)
         y[prefix + (slice(None, 2),)] = 0.
         y[prefix + (-1,)] = 0.
@@ -127,5 +136,7 @@ class LaplaceOperator(EndomorphicOperator):
         ret[sl_l] = deriv
         ret[prefix + (-1,)] = 0.
         ret[sl_r] -= deriv
-        return Field(self.domain, val=from_np(ret)).weight(-1,
-                                                           spaces=self._space)
+        ret = dobj.from_local_data(x.shape, ret, dobj.distaxis(yf))
+        if dobj.distaxis(yf)!= dobj.distaxis(x.val):
+            ret = dobj.redistribute(ret, dist=dobj.distaxis(x.val))
+        return Field(self.domain, val=ret).weight(-1, spaces=self._space)

nifty/operators/power_projection_operator.py

@@ -56,7 +56,7 @@ class PowerProjectionOperator(LinearOperator):
         if dobj.distaxis(x.val) in x.domain.axes[self._space]:  # the distributed axis is part of the projected space
             pindex = dobj.local_data(pindex)
         else:
-            pindex = dobj.to_ndarray(pindex)
+            pindex = dobj.to_global_data(pindex)
         pindex.reshape((1, pindex.size, 1))
         arr = dobj.local_data(x.weight(1).val)
         firstaxis = x.domain.axes[self._space][0]
@@ -74,7 +74,7 @@ class PowerProjectionOperator(LinearOperator):
         if dobj.distaxis(x.val) in x.domain.axes[self._space]:  # the distributed axis is part of the projected space
             pindex = dobj.local_data(pindex)
         else:
-            pindex = dobj.to_ndarray(pindex)
+            pindex = dobj.to_global_data(pindex)
         pindex = pindex.reshape((1, pindex.size, 1))
         arr = dobj.local_data(x.val)
         firstaxis = x.domain.axes[self._space][0]

nifty/spaces/power_space.py

@@ -145,7 +145,7 @@ class PowerSpace(Space):
             locdat = np.searchsorted(tbb, dobj.local_data(k_length_array.val))
             temp_pindex = dobj.from_local_data(
                 k_length_array.val.shape, locdat, dobj.distaxis(k_length_array.val))
-            nbin = len(tbb)
+            nbin = len(tbb)+1
             temp_rho = np.bincount(dobj.local_data(temp_pindex).ravel(),
                                    minlength=nbin)
             temp_rho = dobj.np_allreduce_sum(temp_rho)

nifty/spaces/rg_space.py

@@ -116,7 +116,7 @@ class RGSpace(Space):
             return np.sqrt(np.nonzero(tmp)[0])*self.distances[0]
         else:  # do it the hard way
             # FIXME: this needs to improve for MPI. Maybe unique()/gather()?
-            tmp = np.unique(dobj.to_ndarray(self.get_k_length_array().val))  # expensive!
+            tmp = np.unique(dobj.to_global_data(self.get_k_length_array().val))  # expensive!
             tol = 1e-12*tmp[-1]
             # remove all points that are closer than tol to their right
             # neighbors.

nifty/sugar.py

@@ -224,7 +224,7 @@ def create_power_field(domain, power_spectrum, dtype=None):
     else:
         power_domain = PowerSpace(domain)
         fp = Field(power_domain,
-                   val=dobj.from_ndarray(power_spectrum(power_domain.k_lengths)),
+                   val=dobj.from_global_data(power_spectrum(power_domain.k_lengths)),
                    dtype=dtype)
     P = PowerProjectionOperator(domain, power_domain)
     f = P.adjoint_times(fp)

test/test_minimization/test_minimizers.py

@@ -4,7 +4,6 @@ from numpy.testing import assert_allclose
 import nifty2go as ift
 from itertools import product
 from test.common import expand
-from nifty2go.dobj import to_ndarray as to_np
 
 spaces = [ift.RGSpace([1024], distances=0.123), ift.HPSpace(32)]
 minimizers = [ift.SteepestDescent, ift.RelaxedNewton, ift.VL_BFGS,
@@ -29,6 +28,6 @@ class Test_Minimizers(unittest.TestCase):
 
         (energy, convergence) = minimizer(energy)
         assert convergence == IC.CONVERGED
-        assert_allclose(to_np(energy.position.val),
-                        1./to_np(covariance_diagonal.val),
+        assert_allclose(ift.dobj.to_global_data(energy.position.val),
+                        1./ift.dobj.to_global_data(covariance_diagonal.val),
                         rtol=1e-3, atol=1e-3)

test/test_operators/test_composed_operator.py

@@ -4,7 +4,7 @@ from nifty2go import Field, DiagonalOperator, ComposedOperator
 from test.common import generate_spaces
 from itertools import product
 from test.common import expand
-from nifty2go.dobj import to_ndarray as to_np
+from nifty2go.dobj import to_global_data as to_np
 
 
 class ComposedOperator_Tests(unittest.TestCase):

test/test_operators/test_diagonal_operator.py

@@ -5,7 +5,7 @@ from nifty2go import Field, DiagonalOperator
 from test.common import generate_spaces
 from itertools import product
 from test.common import expand
-from nifty2go.dobj import to_ndarray as to_np
+from nifty2go.dobj import to_global_data as to_np
 
 
 class DiagonalOperator_Tests(unittest.TestCase):

test/test_operators/test_fft_operator.py

@@ -23,7 +23,7 @@ from numpy.testing import assert_allclose
 from nifty2go import Field, RGSpace, LMSpace, HPSpace, GLSpace, FFTOperator
 from itertools import product
 from test.common import expand
-from nifty2go.dobj import to_ndarray as to_np
+from nifty2go import dobj
 
 
 def _get_rtol(tp):
@@ -46,7 +46,7 @@ class FFTOperatorTests(unittest.TestCase):
         inp = Field.from_random(domain=a, random_type='normal', std=7, mean=3,
                                 dtype=itp)
         out = fft.adjoint_times(fft.times(inp))
-        assert_allclose(to_np(inp.val), to_np(out.val), rtol=tol, atol=tol)
+        assert_allclose(dobj.to_global_data(inp.val), dobj.to_global_data(out.val), rtol=tol, atol=tol)
 
     @expand(product([12, 15], [9, 12], [0.1, 1, 3.7],
                     [0.4, 1, 2.7],
@@ -61,7 +61,7 @@ class FFTOperatorTests(unittest.TestCase):
         inp = Field.from_random(domain=a, random_type='normal', std=7, mean=3,
                                 dtype=itp)
         out = fft.adjoint_times(fft.times(inp))
-        assert_allclose(to_np(inp.val), to_np(out.val), rtol=tol, atol=tol)
+        assert_allclose(dobj.to_global_data(inp.val), dobj.to_global_data(out.val), rtol=tol, atol=tol)
 
     @expand(product([0, 1, 2],
                     [np.float64, np.float32, np.complex64, np.complex128]))
@@ -73,7 +73,7 @@ class FFTOperatorTests(unittest.TestCase):
         inp = Field.from_random(domain=(a1, a2, a3), random_type='normal',
                                 std=7, mean=3, dtype=dtype)
         out = fft.adjoint_times(fft.times(inp))
-        assert_allclose(to_np(inp.val), to_np(out.val), rtol=tol, atol=tol)
+        assert_allclose(dobj.to_global_data(inp.val), dobj.to_global_data(out.val), rtol=tol, atol=tol)
 
     @expand(product([0, 3, 6, 11, 30],
                     [np.float64, np.float32, np.complex64, np.complex128]))
@@ -85,7 +85,7 @@ class FFTOperatorTests(unittest.TestCase):
         inp = Field.from_random(domain=a, random_type='normal', std=7, mean=3,
                                 dtype=tp)
         out = fft.adjoint_times(fft.times(inp))
-        assert_allclose(to_np(inp.val), to_np(out.val), rtol=tol, atol=tol)
+        assert_allclose(dobj.to_global_data(inp.val), dobj.to_global_data(out.val), rtol=tol, atol=tol)
 
     @expand(product([128, 256],
                     [np.float64, np.float32, np.complex64, np.complex128]))
@@ -96,7 +96,7 @@ class FFTOperatorTests(unittest.TestCase):
         inp = Field.from_random(domain=a, random_type='normal', std=1, mean=0,
                                 dtype=tp)
         out = fft.adjoint_times(fft.times(inp))
-        assert_allclose(to_np(inp.val), to_np(out.val), rtol=1e-3, atol=1e-1)
+        assert_allclose(dobj.to_global_data(inp.val), dobj.to_global_data(out.val), rtol=1e-3, atol=1e-1)
 
     @expand(product([128, 256],
                     [np.float64, np.float32, np.complex64, np.complex128]))

test/test_spaces/test_lm_space.py

@@ -23,7 +23,7 @@ from numpy.testing import assert_, assert_equal, assert_raises,\
         assert_allclose
 import nifty2go as ift
 from test.common import expand
-from nifty2go.dobj import to_ndarray as to_np
+from nifty2go import dobj
 
 # [lmax, expected]
 CONSTRUCTOR_CONFIGS = [
@@ -93,4 +93,4 @@ class LMSpaceFunctionalityTests(unittest.TestCase):
     @expand(get_k_length_array_configs())
     def test_k_length_array(self, lmax, expected):
         l = ift.LMSpace(lmax)
-        assert_allclose(to_np(l.get_k_length_array().val), expected)
+        assert_allclose(dobj.to_global_data(l.get_k_length_array().val), expected)

test/test_spaces/test_rg_space.py

@@ -22,7 +22,7 @@ import numpy as np
 from numpy.testing import assert_, assert_equal, assert_allclose
 import nifty2go as ift
 from test.common import expand
-from nifty2go.dobj import to_ndarray as to_np
+from nifty2go import dobj
 
 # [shape, distances, harmonic, expected]
 CONSTRUCTOR_CONFIGS = [
@@ -111,7 +111,7 @@ class RGSpaceFunctionalityTests(unittest.TestCase):
     @expand(get_k_length_array_configs())
     def test_k_length_array(self, shape, distances, expected):
         r = ift.RGSpace(shape=shape, distances=distances, harmonic=True)
-        assert_allclose(to_np(r.get_k_length_array().val), expected)
+        assert_allclose(dobj.to_global_data(r.get_k_length_array().val), expected)
 
     @expand(get_dvol_configs())
     def test_dvol(self, shape, distances, harmonic, power):