ML with separate gradient steps

ptypy/accelerate/base/kernels.py

@@ -269,6 +269,37 @@ def make_model(self, b_aux, addr):
         tf = aux.reshape(sh[0], self.nmodes, sh[1], sh[2])
         Imodel[:] = ((tf * tf.conj()).real).sum(1)
 
+    def make_a012_notb(self, b_f, b_a, addr, I, fic):
+
+        # reference shape (= GPU global dims)
+        sh = I.shape
+
+        # stopper
+        maxz = I.shape[0]
+
+        A0 = self.npy.Imodel
+        A1 = self.npy.LLerr
+        A2 = self.npy.LLden
+
+        # batch buffers
+        f = b_f[:maxz * self.nmodes]
+        a = b_a[:maxz * self.nmodes]
+
+        ## Actual math ## (subset of FUK.fourier_error)
+        fc = fic.reshape((maxz,1,1))
+        A0.fill(0.)
+        tf = np.real(f * f.conj()).astype(self.ftype)
+        A0[:maxz] = np.double(tf.reshape(maxz, self.nmodes, sh[1], sh[2]).sum(1) * fc) - I
+
+        A1.fill(0.)
+        tf = 2. * np.real(f * a.conj())
+        A1[:maxz] = tf.reshape(maxz, self.nmodes, sh[1], sh[2]).sum(1) * fc
+
+        A2.fill(0.)
+        tf = np.real(a * a.conj())
+        A2[:maxz] = tf.reshape(maxz, self.nmodes, sh[1], sh[2]).sum(1) * fc
+        return
+
     def make_a012(self, b_f, b_a, b_b, addr, I, fic):
 
         # reference shape (= GPU global dims)

ptypy/accelerate/cuda_common/make_a012_notb.cu

@@ -0,0 +1,61 @@
+/** fmag_all_update without b.
+ *
+ * Data types:
+ * - IN_TYPE: the data type for the inputs (float or double)
+ * - OUT_TYPE: the data type for the outputs (float or double)
+ * - MATH_TYPE: the data type used for computation 
+ * - ACC_TYPE: data type used for accumulation
+ */
+
+#include "common.cuh"
+
+extern "C" __global__ void make_a012_notb(const complex<IN_TYPE>* f,
+                                          const complex<IN_TYPE>* a,
+                                          const IN_TYPE* I,
+                                          const IN_TYPE* fic,
+                                          OUT_TYPE* A0,
+                                          OUT_TYPE* A1,
+                                          OUT_TYPE* A2,
+                                          int z,
+                                          int y,
+                                          int x,
+                                          int maxz)
+{
+  int ix = threadIdx.x + blockIdx.x * blockDim.x;
+  int iz = blockIdx.z;
+
+  if (ix >= x)
+    return;
+
+  if (iz >= maxz)
+  {
+    A0[iz * x + ix] = OUT_TYPE(0);  // make sure it's the right type (double/float)
+    A1[iz * x + ix] = OUT_TYPE(0);
+    A2[iz * x + ix] = OUT_TYPE(0);
+    return;
+  }
+
+  // we sum across y directly, as this is the number of modes,
+  // which is typically small
+  auto sumtf0 = ACC_TYPE(0);
+  auto sumtf1 = ACC_TYPE(0);
+  auto sumtf2 = ACC_TYPE(0);
+  for (auto iy = 0; iy < y; ++iy)
+  {
+    complex<MATH_TYPE> fv = f[iz * y * x + iy * x + ix];
+    sumtf0 += fv.real() * fv.real() + fv.imag() * fv.imag();
+
+    complex<MATH_TYPE> av = a[iz * y * x + iy * x + ix];
+    // 2 * real(f * conj(a))
+    sumtf1 += MATH_TYPE(2) * (fv.real() * av.real() + fv.imag() * av.imag());
+
+    // abs(a)^2
+    sumtf2 += av.real() * av.real() + av.imag() * av.imag();
+  }
+
+  MATH_TYPE Iv = I[iz * x + ix];
+  MATH_TYPE ficv = fic[iz];
+  A0[iz * x + ix] = OUT_TYPE(MATH_TYPE(sumtf0) * ficv - Iv);
+  A1[iz * x + ix] = OUT_TYPE(MATH_TYPE(sumtf1) * ficv);
+  A2[iz * x + ix] = OUT_TYPE(MATH_TYPE(sumtf2) * ficv);
+}

ptypy/accelerate/cuda_cupy/kernels.py

@@ -676,6 +676,7 @@ def __init__(self, aux, nmodes=1, queue=None, accumulate_type='double', math_typ
         }
         self.make_model_cuda = load_kernel('make_model', subs)
         self.make_a012_cuda = load_kernel('make_a012', subs)
+        self.make_a012_notb_cuda = load_kernel('make_a012_notb', subs)
         self.error_reduce_cuda = load_kernel('error_reduce', {
             **subs,
             'OUT_TYPE': 'float' if self.ftype == np.float32 else 'double',
@@ -748,6 +749,29 @@ def make_a012(self, b_f, b_a, b_b, addr, I, fic):
                             args=(b_f, b_a, b_b, I, fic,
                             A0, A1, A2, z, y, x, maxz))
 
+    def make_a012_notb(self, b_f, b_a, addr, I, fic):
+        # reference shape (= GPU global dims)
+        sh = I.shape
+
+        # stopper
+        maxz = I.shape[0]
+
+        A0 = self.gpu.Imodel
+        A1 = self.gpu.LLerr
+        A2 = self.gpu.LLden
+
+        z = np.int32(sh[0])
+        maxz = np.int32(maxz)
+        y = np.int32(self.nmodes)
+        x = np.int32(sh[1]*sh[2])
+        bx = 1024
+        if self.queue is not None:
+            self.queue.use()
+        self.make_a012_notb_cuda(grid=(int((x + bx - 1) // bx), 1, int(z)),
+                            block=(bx, 1, 1),
+                            args=(b_f, b_a, I, fic,
+                            A0, A1, A2, z, y, x, maxz))
+
     def fill_b(self, addr, Brenorm, w, B):
         # stopper
         maxz = w.shape[0]

ptypy/accelerate/cuda_pycuda/kernels.py

@@ -654,6 +654,7 @@ def __init__(self, aux, nmodes=1, queue=None, accumulate_type = 'double', math_t
         }
         self.make_model_cuda = load_kernel('make_model', subs)
         self.make_a012_cuda = load_kernel('make_a012', subs)
+        self.make_a012_notb_cuda = load_kernel('make_a012_notb', subs)
         self.error_reduce_cuda = load_kernel('error_reduce', {
             **subs,
             'OUT_TYPE': 'float' if self.ftype == np.float32 else 'double',
@@ -723,6 +724,28 @@ def make_a012(self, b_f, b_a, b_b, addr, I, fic):
                             grid=(int((x + bx - 1) // bx), 1, int(z)),
                             stream=self.queue)
 
+    def make_a012_notb(self, b_f, b_a, addr, I, fic):
+        # reference shape (= GPU global dims)
+        sh = I.shape
+
+        # stopper
+        maxz = I.shape[0]
+
+        A0 = self.gpu.Imodel
+        A1 = self.gpu.LLerr
+        A2 = self.gpu.LLden
+
+        z = np.int32(sh[0])
+        maxz = np.int32(maxz)
+        y = np.int32(self.nmodes)
+        x = np.int32(sh[1]*sh[2])
+        bx = 1024
+        self.make_a012_notb_cuda(b_f, b_a, I, fic,
+                            A0, A1, A2, z, y, x, maxz,
+                            block=(bx, 1, 1),
+                            grid=(int((x + bx - 1) // bx), 1, int(z)),
+                            stream=self.queue)
+
     def fill_b(self, addr, Brenorm, w, B):
         # stopper
         maxz = w.shape[0]