revert numba updates in pcg

ratt-ru · Sep 23, 2024 · 55882b6 · 55882b6
1 parent 90b85db
commit 55882b6
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 18 deletions.
diff --git a/pfb/__init__.py b/pfb/__init__.py
@@ -11,11 +11,11 @@
 __version__ = '0.0.4'
 
 def set_envs(nthreads, ncpu):
-    os.environ["OMP_NUM_THREADS"] = '2'
-    os.environ["OPENBLAS_NUM_THREADS"] = '2'
-    os.environ["MKL_NUM_THREADS"] = '2'
-    os.environ["VECLIB_MAXIMUM_THREADS"] = '2'
-    os.environ["NPY_NUM_THREADS"] = '2'
+    os.environ["OMP_NUM_THREADS"] = str(nthreads)
+    os.environ["OPENBLAS_NUM_THREADS"] = str(nthreads)
+    os.environ["MKL_NUM_THREADS"] = str(nthreads)
+    os.environ["VECLIB_MAXIMUM_THREADS"] = str(nthreads)
+    os.environ["NPY_NUM_THREADS"] = str(nthreads)
     os.environ["NUMBA_NUM_THREADS"] = str(nthreads)
     os.environ["JAX_PLATFORMS"] = 'cpu'
     os.environ["JAX_ENABLE_X64"] = 'True'

diff --git a/pfb/opt/pcg.py b/pfb/opt/pcg.py
@@ -223,15 +223,29 @@ def M(x): return x
         Ap = A(p)
         tA += (time() - ti)
         ti = time()
-        # rnorm = np.vdot(r, y)
-        # alpha = rnorm / np.vdot(p, Ap)
+        rnorm = np.vdot(r, y)
+        alpha = rnorm / np.vdot(p, Ap)
         # import ipdb; ipdb.set_trace()
-        rnorm, alpha = alpha_update(r, y, p, Ap)
+        # rnorm, alpha = alpha_update(r, y, p, Ap)
         tvdot += (time() - ti)
         ti = time()
         # x = xp + alpha * p
         # r = rp + alpha * Ap
-        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)
 
@@ -243,12 +257,20 @@ def M(x): return x
         #     rnorm_next = np.vdot(r, y)
 
         ti = time()
-        # rnorm_next = np.vdot(r, y)
-        # beta = rnorm_next / rnorm
+        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')
+
         # p = beta * p - y
-        rnorm, p = beta_update(r, y, p, rnorm)
+        # rnorm, p = beta_update(r, y, p, rnorm)
         tp += (time() - ti)
-        # rnorm = rnorm_next
+        rnorm = rnorm_next
         k += 1
         epsp = eps
         ti = time()

diff --git a/pfb/workers/model2comps.py b/pfb/workers/model2comps.py
@@ -336,8 +336,3 @@ def _model2comps(**kw):
                   fits_name,
                   hdr,
                   overwrite=True)
-
-
-
-
-    print("All done here.", file=log)