JP-3653: Nsclean speedup

CHANGES.rst

@@ -138,6 +138,13 @@ master_background_mos
   ``slit.source_name`` now contains the string "BKG" instead of
   "background". [#8533]
 
+nsclean
+-------
+
+- Improved run time of NSClean.fit() by using a vector rather than a large,
+  sparse matrix to perform linear algebra operations with a diagonal weight
+  matrix. [#8547]
+
 outlier_detection
 -----------------
 

jwst/nsclean/lib.py

@@ -134,7 +134,10 @@
 
             # Get data and weights for this line
             d = data[y][self.mask[y]]  # unmasked (useable) data
-            p = np.diag(self.P[y][self.mask[y]])  # Weights
+            # The line below uses a vector to represent a diagonal weight matrix.
+            # Multiplications by this vector later on may be viewed as
+            # equivalent formulations to multiplication by diag(p).
+            p = self.P[y][self.mask[y]]  # Weights
 
             # If none of the pixels in this line is useable (all masked out),
             # skip and move on to the next line.
@@ -164,15 +167,15 @@
 
             # Compute the Moore-Penrose inverse of A = P*B.
             #     $A^+ = (A^H A)^{-1} A^H$
-            A = np.matmul(p, B)
+            A = B*p[:, np.newaxis]
             AH = np.conjugate(A.transpose())  # Hermitian transpose of A
             pinv_PB = np.matmul(np.linalg.inv(np.matmul(AH, A)), AH)
 
             # Solve for the Fourier transform of this line's background samples.
             # The way that we have done it, this multiplies the input data by the 
             # number of samples used for the fit.
             rfft = np.zeros(self.nx//2 + 1, dtype=np.complex64)
-            rfft[:k.shape[1]] = np.matmul(np.matmul(pinv_PB, p), d)
+            rfft[:k.shape[1]] = np.matmul(pinv_PB, p*d)
 
             # Numpy requires that the forward transform multiply
             # the data by n. Correct normalization.