resize ducc thread pool in comps2vis

ratt-ru · Sep 26, 2024 · df9c779 · df9c779
1 parent 1a61ca9
commit df9c779
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 27 deletions.
diff --git a/pfb/operators/gridder.py b/pfb/operators/gridder.py
@@ -243,6 +243,9 @@ def _comps2vis_impl(uvw,
                     ncorr_out=4,
                     freq_min=-np.inf,
                     freq_max=np.inf):
+    # why is this necessary?
+    resize_thread_pool(nthreads)
+
     # adjust for chunking
     # need a copy here if using multiple row chunks
     rbin_idx2 = rbin_idx - rbin_idx.min()

diff --git a/pfb/opt/pcg.py b/pfb/opt/pcg.py
@@ -186,27 +186,25 @@ def M(x): return x
         ti = time()
         rnorm = np.vdot(r, y)
         alpha = rnorm / np.vdot(p, Ap)
-        # import ipdb; ipdb.set_trace()
-        # rnorm, alpha = alpha_update(r, y, p, Ap)
         tvdot += (time() - ti)
         ti = time()
         # x = xp + alpha * p
         # r = rp + alpha * Ap
-        # ne.evaluate('xp + alpha*p',
-        #             out=x,
-        #             local_dict={
-        #                 'xp': xp,
-        #                 'alpha': alpha,
-        #                 'p': p},
-        #             casting='unsafe')
-        # ne.evaluate('rp + alpha*Ap',
-        #             out=r,
-        #             local_dict={
-        #                 'rp': rp,
-        #                 'alpha': alpha,
-        #                 'Ap': Ap},
-        #             casting='unsafe')
-        x, r = update(x, xp, r, rp, p, Ap, alpha)
+        ne.evaluate('xp + alpha*p',
+                    out=x,
+                    local_dict={
+                        'xp': xp,
+                        'alpha': alpha,
+                        'p': p},
+                    casting='unsafe')
+        ne.evaluate('rp + alpha*Ap',
+                    out=r,
+                    local_dict={
+                        'rp': rp,
+                        'alpha': alpha,
+                        'Ap': Ap},
+                    casting='unsafe')
+        # x, r = update(x, xp, r, rp, p, Ap, alpha)
         tupdate += (time() - ti)
         y = M(r)
 
@@ -220,17 +218,16 @@ def M(x): return x
         ti = time()
         rnorm_next = np.vdot(r, y)
         beta = rnorm_next / rnorm
-        # ne.evaluate('beta*p-y',
-        #             out=p,
-        #             local_dict={
-        #                 'beta': beta,
-        #                 'p': p,
-        #                 'y': y},
-        #             casting='unsafe')
-
+        ne.evaluate('beta*p-y',
+                    out=p,
+                    local_dict={
+                        'beta': beta,
+                        'p': p,
+                        'y': y},
+                    casting='unsafe')
         # p = beta * p - y
-        p *= beta
-        p -= y
+        # p *= beta
+        # p -= y
         # rnorm, p = beta_update(r, y, p, rnorm)
         tp += (time() - ti)
         rnorm = rnorm_next