weights and mask

desihub · Apr 10, 2016 · a8b2444 · a8b2444
1 parent fcce365
commit a8b2444
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/py/desiutil/funcfits.py b/py/desiutil/funcfits.py
@@ -167,8 +167,12 @@ def iter_fit(xarray, yarray, func, order, weights=None, sigma=None,
         w = np.where(mask == 0)
         xfit = xarray[w]
         yfit = yarray[w]
+        if weights is not None:
+            wfit = weights[w]
+        else:
+            wfit = None
         # Fit
-        dfit = func_fit(xfit, yfit, func, order, xmin=xmin, xmax=xmax,
+        dfit = func_fit(xfit, yfit, func, order, xmin=xmin, xmax=xmax, w=wfit,
                         **kwargs)
         yrng = func_val(xarray, dfit)
         # Reject
@@ -208,8 +212,8 @@ def iter_fit(xarray, yarray, func, order, weights=None, sigma=None,
             if mskcnt-imskcnt > max_rej:
                 break
         mskcnt = np.sum(mask)
-        w = np.where(mask == 0)
     # Final fit
+    w = np.where(mask == 0)
     xfit = xarray[w]
     yfit = yarray[w]
     fdict = func_fit(xfit, yfit, func, order, xmin=xmin, xmax=xmax, **kwargs)

diff --git a/py/desiutil/test/test_funcfits.py b/py/desiutil/test/test_funcfits.py
@@ -111,7 +111,7 @@ def test_iterfit(self):
         y2 = func_val(x2, dfit)
         np.testing.assert_allclose(y2[50], 0.99941444872371643)
 
-    def test_iterfit(self):
+    def test_iterfit2(self):
         """Test iter fit with some special cases.
         """
         # Generate data