Repository: kymatio/kymatio

Issue:
3D output format
Currently, `Scattering3D` takes an input of shape `(B, M, N, O)`, where `B` is the batch size and `M, N, O` is the shape of the input signal. Depending on `method`, the output takes different formats:

- If `method='standard'`, the output is of shape `(B, 2*M*N*O, total_scales, L+1)`, where `total_scales` is the number of scales at all orders, that is `(J+1)+J*(J+1)/2` if `order_2=True`. The factor `2` in the second axis comes from the fact that these are still stored as "complex" values (even though they are real), so every other element along that axis is zero. These zeroes should probably be eliminated.
- If `method='integral'`, the output is of shape `(B, len(integral_powers), total_scales, L+1)`, where `integral_powers` is the argument provided to `forward()`.
- If `method='local'`, the output is of shape `(B, points.shape[1], total_scales, L+1)`, where `points` is the argument provided to `forward()`.

The first question is whether we should provide all these outputs through the same interface. To a certain extent, `'integral'` and `'local'` can be computed from `'standard'`, but this may not be computationally optimal.

The second question is how to structure these outputs to be more consistent with 1D and 2D. I suggest that we adopt a similar output shape `(B, C, M, N, O)`, where the second axis comingles orders, scales, ells, integral powers, and points, as the case may be. Alternatively, for `'integral`' and `'local'`, the third axis could index `integral_powers` or `points` instead of a singleton spatial dimension.


Discussion:
Need to figure this out before release... @eickenberg, @louity?
> If method='standard', the output is of shape (B, 2*M*N*O, total_scales, L+1), where total_scales is the number of scales at all orders, that is (J+1)+J*(J+1)/2 if order_2=True.

😱 😱 ❗️ ❗️ ❗️ 

Why do we ravel space and not keep it like every other part of this software package?

> These zeroes should probably be eliminated.

yeah

> where integral_powers is the argument provided to forward()
👎 

> where points is the argument provided to forward()

maybe OK.

> 'integral' and 'local' can be computed from 'standard', but this may not be computationally optimal.

indeed, though I do not know whether the code actually does anything to optimize. But I think these possibilities should be kept because 3D can easily fill up GPU memory.

> I suggest that we adopt a similar output shape (B, C, M, N, O), where the second axis comingles orders, scales, ells, integral powers, and points, as the case may be.

yeah I agree

> Alternatively, for 'integral' and 'local', the third axis could index integral_powers or points instead of a singleton spatial dimension.

yeah why not


I agree that this needs to be addressed. 
> Why do we ravel space and not keep it like every other part of this software package?

:man_shrugging: 
 
> > These zeroes should probably be eliminated.
> 
> yeah

I'm glad we agree on this.
 
> > where integral_powers is the argument provided to forward()
>
> -1

So make this an attribute? This is what #230 suggests.
 
> > where points is the argument provided to forward()
> 
> maybe OK.

Again, #230 is the place to discuss this. For consistency, I'd prefer no arguments to `forward` except for the array.

> > 'integral' and 'local' can be computed from 'standard', but this may not be computationally optimal.
> 
> indeed, though I do not know whether the code actually does anything to optimize. But I think these possibilities should be kept because 3D can easily fill up GPU memory.

Well, you skip a lowpass filter in `integral` and I think `local` replaces that lowpass convolution with an inner product, so those should be non-trivial speedups.

> > I suggest that we adopt a similar output shape (B, C, M, N, O), where the second axis comingles orders, scales, ells, integral powers, and points, as the case may be.
> 
> yeah I agree

Ok good.

> 
> > Alternatively, for 'integral' and 'local', the third axis could index integral_powers or points instead of a singleton spatial dimension.
> 
> yeah why not

Good.

> I agree that this needs to be addressed.

Ok. So if I've understood correctly, this is the plan:

- Remove extraneous zeros in `standard`.
- Do not ravel space in `standard`.
- Make third axis index `integral_powers` (for `integral`) or `points` (for `local`).
- (Remove kwargs and put them as attributes, but again, this is #230)

Does that sound correct?
while we're at it, i should add that I don't find the kwarg name `method` to be informative enough. I had this discussion with @eickenberg last month. Him and I came up with the name `sampling` as a replacement for `method`.

Resolving patch:
diff --git a/kymatio/scattering3d/scattering3d.py b/kymatio/scattering3d/scattering3d.py
--- a/kymatio/scattering3d/scattering3d.py
+++ b/kymatio/scattering3d/scattering3d.py
@@ -118,8 +118,7 @@ def _compute_standard_scattering_coefs(self, input_array):
 
         """
         convolved_input = self._low_pass_filter(input_array, self.J)
-        return subsample(convolved_input, self.J).view(
-            input_array.size(0), -1, 1)
+        return subsample(convolved_input, self.J)
 
     def _compute_local_scattering_coefs(self, input_array, points, j):
         """
@@ -354,16 +353,15 @@ def forward(self, input_array, max_order=2, rotation_covariant=True,
                         conv_modulus, l, j_2)
                     s_order_2_l.append(compute_scattering_coefs(
                         conv_modulus_2, method, method_args, j_2))
-            s_order_1.append(torch.cat(s_order_1_l, -1))
+            s_order_1.append(torch.stack([arr[..., 0] for arr in s_order_1_l], 1))
             if max_order == 2:
-                s_order_2.append(torch.cat(s_order_2_l, -1))
-
+                s_order_2.append(torch.stack([arr[..., 0] for arr in s_order_2_l], 1))
         if max_order == 2:
             return torch.cat(
-                [torch.stack(s_order_1, dim=-1),
-                torch.stack(s_order_2, dim=-1)], -2)
+                [torch.stack(s_order_1, dim=2),
+                torch.stack(s_order_2, dim=2)], 1)
         else:
-            return torch.stack(s_order_1, dim=-1)
+            return torch.stack(s_order_1, dim=2)
 
 
     __call__ = forward
diff --git a/kymatio/scattering3d/utils.py b/kymatio/scattering3d/utils.py
--- a/kymatio/scattering3d/utils.py
+++ b/kymatio/scattering3d/utils.py
@@ -52,9 +52,7 @@ def generate_weighted_sum_of_gaussians(grid, positions, weights, sigma,
 
 
 def subsample(input_array, j):
-    return input_array.unfold(3, 1, 2 ** j
-                     ).unfold(2, 1, 2 ** j
-                     ).unfold(1, 1, 2 ** j).contiguous()
+    return input_array[..., ::2 ** j, ::2 ** j, ::2 ** j, :].contiguous()
 
 
 


Test patch:
diff --git a/kymatio/scattering3d/tests/test_scattering3d.py b/kymatio/scattering3d/tests/test_scattering3d.py
--- a/kymatio/scattering3d/tests/test_scattering3d.py
+++ b/kymatio/scattering3d/tests/test_scattering3d.py
@@ -176,8 +176,14 @@ def test_against_standard_computations():
 
         # Extract orders and make order axis the slowest in accordance with
         # the stored reference scattering transform.
-        order_1 = orders_1_and_2[:,:,0:n_order_1,:]
-        order_2 = orders_1_and_2[:,:,n_order_1:n_order_1+n_order_2,:]
+        order_1 = orders_1_and_2[:,0:n_order_1,...]
+        order_2 = orders_1_and_2[:,n_order_1:n_order_1+n_order_2,...]
+
+        # Permute the axes since reference has (batch index, integral power, j,
+        # ell) while the computed transform has (batch index, j, ell, integral
+        # power).
+        order_1 = order_1.permute(0, 3, 1, 2)
+        order_2 = order_2.permute(0, 3, 1, 2)
 
         order_1 = order_1.reshape((batch_size, -1))
         order_2 = order_2.reshape((batch_size, -1))
@@ -228,7 +234,7 @@ def test_solid_harmonic_scattering():
             sigma_wavelet = sigma_0_wavelet*2**j
             k = sigma_wavelet / np.sqrt(sigma_wavelet**2 + sigma_gaussian**2)
             for l in range(1, L+1):
-                err = torch.abs(s[0, 0, j, l] - k ** l).sum()/(1e-6+s[0, 0, j, l].abs().sum())
+                err = torch.abs(s[0, j, l, 0] - k ** l).sum()/(1e-6+s[0, j, l, 0].abs().sum())
                 assert err<1e-4
 
 def test_larger_scales():