ENH: add newaxis to TensorSpace and use it in DiscreteLp.astype

odl/discr/lp_discr.py

@@ -473,7 +473,8 @@ def _astype(self, dtype):
             tspace = self.tspace.astype(dtype)
             axis_labels_new = self.axis_labels
         else:
-            tspace = self.tspace.byaxis[self.ndim_out:].astype(dtype)
+            tspace = self.tspace.byaxis[self.ndim_out:].newaxis(0, dtype.shape)
+            tspace = tspace.astype(dtype.base)
             axis_labels_extra = tuple('$i_{}$'.format(ax)
                                       for ax in range(len(dtype.shape)))
             axis_labels_new = (axis_labels_extra +

odl/space/base_tensors.py

@@ -270,11 +270,28 @@ def astype(self, dtype):
         else:
             return self._astype(dtype)
 
-    @property
-    def default_order(self):
-        """Default storage order for new elements in this space.
+    def newaxis(self, index, shape, weighting=None):
+        """Return a copy of this space with extra axes inserted.
 
-        This property should be overridden by subclasses.
+        .. note::
+            This function is most useful in the default case of per-axis
+            weighting factors. For weighting with a single constant
+            or with an array of the same shape as this space, "inserting"
+            new factors does not really make sense, and the ``weighting``
+            argument will not be used. In these cases, it is better to
+            construct a new space altogether.
+
+        Parameters
+        ----------
+        index : int
+            Position at which to insert new axes. Negative indices are
+            counted backwards from the end.
+        shape : nonnegative int or sequence of nonnegative ints
+            Sizes of the axes that are to be inserted.
+        weighting : float, `array-like` or `Weighting`, optional
+            Weighting factors that should be used in the new axes.
+            The number of factors must correspond to the number of entries
+            in ``shape`` (1 if ``shape`` is an int).
         """
         raise NotImplementedError('abstract method')
 
@@ -483,6 +500,14 @@ def _divide(self, x1, x2, out):
         """
         raise NotImplementedError('abstract method')
 
+    @property
+    def default_order(self):
+        """Default storage order for new elements in this space.
+
+        This property should be overridden by subclasses.
+        """
+        raise NotImplementedError('abstract method')
+
     @staticmethod
     def default_dtype(field=None):
         """Return the default data type for a given field.

odl/space/npy_tensors.py

@@ -92,8 +92,8 @@ def __init__(self, shape, dtype=None, **kwargs):
             Exponent of the norm. For values other than 2.0, no
             inner product is defined.
 
-            This option has no impact if either ``dist``, ``norm`` or
-            ``inner`` is given, or if ``dtype`` is non-numeric.
+            This option has no impact if either of ``weighting``, ``dist``,
+            ``norm`` or ``inner`` is given, or if ``dtype`` is non-numeric.
 
             Default: 2.0
 
@@ -103,7 +103,8 @@ def __init__(self, shape, dtype=None, **kwargs):
             Use weighted inner product, norm, and dist. The following
             types are supported as ``weighting``:
 
-            - ``None``: No weighting, i.e. weighting with ``1.0`` (default).
+            - ``None``: No weighting, i.e. weighting with ``1.0`` for
+              each axis (default).
             - ``float``: Weighting by a constant.
             - sequence of length ``ndim``: Separate weighting per axis.
               Entries can be constants or 1D arrays.
@@ -312,8 +313,9 @@ def __init__(self, shape, dtype=None, **kwargs):
         elif inner is not None:
             self.__weighting = NumpyTensorSpaceCustomInner(inner)
         else:
-            # No weighting, i.e., weighting with constant 1.0
-            self.__weighting = NumpyTensorSpaceConstWeighting(1.0, exponent)
+            # No weighting, i.e., weighting with constant 1.0 in each axis
+            self.__weighting = NumpyTensorSpacePerAxisWeighting(
+                [1.0] * self.ndim, exponent)
 
         # Make sure there are no leftover kwargs
         if kwargs:
@@ -336,10 +338,11 @@ def weighting(self):
 
     @property
     def is_weighted(self):
-        """Return ``True`` if the space is not weighted by constant 1.0."""
+        """Return ``True`` if the space is not weighted by 1.0 in each axis."""
         return not (
-            isinstance(self.weighting, NumpyTensorSpaceConstWeighting) and
-            self.weighting.const == 1.0)
+            isinstance(self.weighting, NumpyTensorSpacePerAxisWeighting) and
+            self.weighting.array_axes == () and
+            self.weighting.consts == (1.0,) * self.ndim)
 
     @property
     def exponent(self):
@@ -478,6 +481,8 @@ def _astype(self, dtype):
         """
         kwargs = {}
         dtype = np.dtype(dtype)
+        if dtype.shape != ():
+            raise ValueError('`dtype` with shape not allowed')
 
         # Use weighting only for floating-point types, otherwise, e.g.,
         # `space.astype(bool)` would fail
@@ -498,6 +503,107 @@ def _astype(self, dtype):
 
         return type(self)(dtype.shape + self.shape, dtype=dtype.base, **kwargs)
 
+    def newaxis(self, index, shape, weighting=None):
+        """Return a copy of this space with extra axes inserted.
+
+        .. note::
+            This function is most useful in the default case of per-axis
+            weighting factors. For weighting with a single constant, with
+            an array of the same shape as this space, or for custom
+            weighting classes, "inserting" new factors does not really make
+            sense. This function will therefore use the old constant in the
+            first case and fall back to no weighting in all other non-standard
+            cases. If that is not the desired outcome, a new space should
+            be constructed altogether.
+
+        Parameters
+        ----------
+        index : int
+            Position at which to insert new axes. Negative indices are
+            counted backwards from the end.
+        shape : nonnegative int or sequence of nonnegative ints
+            Sizes of the axes that are to be inserted.
+        weighting : `array-like` or sequence of `array-like`, optional
+            Weighting factors that should be used in the new axes.
+            The number of factors must correspond to the number of entries
+            in ``shape`` (1 if ``shape`` is an int).
+
+        Examples
+        --------
+        In the default unweighted case, the extended space is unweighted
+        as well, unless specified otherwise:
+
+        >>> space = odl.rn(3)
+        >>> space.newaxis(0, 2)
+        rn((2, 3))
+        >>> space.newaxis(1, 2)
+        rn((3, 2))
+        >>> space = odl.rn((2, 5))
+        >>> space.newaxis(1, (3, 4))
+        rn((2, 3, 4, 5))
+        >>> space.newaxis(1, (3, 4), weighting=[0.5, 2.0])
+        rn((2, 3, 4, 5), weighting=(1.0, 0.5, 2.0, 1.0))
+
+        The (default) per-axis weighting is extended with factors 1.0 by
+        default, or other factors as given:
+
+        >>> space = odl.rn((2, 4))
+        >>> space.newaxis(1, 3) == space.newaxis(1, 3, weighting=[1.0])
+        True
+        >>> space.newaxis(1, 3, weighting=[2, 1, 0.5])
+        rn((2, 3, 4), weighting=(1.0, [ 2. ,  1. ,  0.5], 1.0))
+        """
+        # Check inputs
+        index, index_in = int(index), index
+        if index != index_in:
+            raise TypeError('`index` must be an int, got {!r}'
+                            ''.format(index_in))
+        if index < 0:
+            index += self.ndim
+        if not 0 <= index <= self.ndim:
+            raise IndexError(
+                '`index` must satisfy `-{ndim} <= index <= {ndim}`, got {0}'
+                ''.format(index_in, ndim=self.ndim))
+
+        try:
+            shape = tuple(shape)
+        except TypeError:
+            shape = (shape,)
+
+        num_axes = len(shape)
+
+        # Determine new weighting
+        if isinstance(self.weighting, PerAxisWeighting):
+            if weighting is None:
+                weighting = [1.0] * num_axes
+            else:
+                if num_axes == 1:
+                    weighting = [
+                        np.asarray(weighting, dtype=self.dtype).squeeze()
+                    ]
+                else:
+                    weighting = [np.asarray(w, dtype=self.dtype)
+                                 for w in weighting]
+
+            if len(weighting) != num_axes:
+                raise ValueError(
+                    '`shape` has length {}, but got {} weighting factors'
+                    ''.format(num_axes, len(weighting)))
+
+            old_factors = list(self.weighting.factors)
+            new_factors = old_factors[:index] + weighting + old_factors[index:]
+            new_weighting = NumpyTensorSpacePerAxisWeighting(
+                new_factors, self.exponent)
+
+        elif isinstance(self.weighting, ConstWeighting):
+            new_weighting = self.weighting
+        else:
+            new_weighting = None
+
+        # Construct new space
+        new_shape = self.shape[:index] + shape + self.shape[index:]
+        return type(self)(new_shape, self.dtype, weighting=new_weighting)
+
     def zero(self):
         """Return a tensor of all zeros.
 
@@ -2588,7 +2694,7 @@ def __eq__(self, other):
 
         same_const_idcs = (other.const_axes == self.const_axes)
         consts_equal = (self.consts == other.consts)
-        arrs_ident = (a is b for a, b in zip(self.arrays, other.arrays))
+        arrs_ident = all(a is b for a, b in zip(self.arrays, other.arrays))
 
         return (super(NumpyTensorSpacePerAxisWeighting, self).__eq__(other) and
                 same_const_idcs and

odl/space/weighting.py

@@ -561,7 +561,7 @@ def __hash__(self):
     @property
     def repr_part(self):
         """String usable in a space's ``__repr__`` method."""
-        optargs = [('weighting', self.const, 1.0),
+        optargs = [('weighting', self.const, None),
                    ('exponent', self.exponent, 2.0)]
         return signature_string([], optargs, mod=':.4')