Add dict option for max_order to allow finer control of which cross terms to include

Author: rmjarvis

piff/basis_interp.py

@@ -463,8 +463,13 @@ class BasisPolynomial(BasisInterp):
                         [default: ('u','v')]
     :param max_order:   The maximum total order to use for cross terms between keys.
                         [default: None, which uses the maximum value of any individual key's order]
+                        If this is an integer, it applies to all pairs, but you may also specify
+                        a dict mapping pairs of keys to an integer.  E.g. {('u','v'):3,
+                        ('u','z'):0, ('v','z'):0}.  This sets the maximum order for cross terms
+                        between these pairs.  Furthermore, any pairs for which you want to skip
+                        cross terms (max=0) may be omitted from the dict.
     :param solver:      Which solver to use.  Solvers available are "scipy", "qr", "jax",
-                        "cpp". See above for details.
+                        "cpp". See above for details. [default: 'scipy']
     :param logger:      A logger object for logging debug info. [default: None]
     """
     _type_name = 'BasisPolynomial'
@@ -515,12 +520,13 @@ def __init__(
             logger.warning("JAX not installed. Reverting to numpy/scipy.")
             self.solver = "scipy"
 
-        if self._max_order<0 or np.any(np.array(self._orders) < 0):
+        if np.any(np.array(self._orders) < 0):
             # Exception if we have any requests for negative orders
             raise ValueError('Negative polynomial order specified')
 
         self.kwargs = {
             'order' : order,
+            'max_order' : max_order,
             'keys' : keys,
             'solver': solver,
         }
@@ -529,9 +535,38 @@ def __init__(
         # Start with 1d arrays giving orders in all dimensions
         ord_ranges = [np.arange(order+1,dtype=int) for order in self._orders]
         # Nifty trick to produce n-dim array holding total order
-        #sumorder = np.sum(np.ix_(*ord_ranges))  # This version doesn't work in numpy 1.19
         sumorder = np.sum(np.meshgrid(*ord_ranges, indexing='ij'), axis=0)
-        self._mask = sumorder <= self._max_order
+
+        if isinstance(self._max_order, dict):
+            # This code is not particularly efficient.  Hopefully it doesn't matter.
+            # Basically set a maxorder for each element in sumorder based on whether it is
+            # a) a power of a single key.  Use the order for that key.
+            # b) a cross-product of multiple keys.  Use it only if it is in the max_order dict.
+            max_orders = np.zeros_like(sumorder)
+
+            def get_indices(arr, pre=()):
+                # Get the index tuples of the given multi-dimensional array.
+                if not isinstance(arr, np.ndarray):
+                    yield pre
+                else:
+                    for i in range(len(arr)):
+                        yield from get_indices(arr[i], pre + (i,))
+            for index in get_indices(sumorder):
+                for k, order in enumerate(self._orders):
+                    if index[k] > 0 and all(index[j] == 0 for j in range(len(index)) if j != k):
+                        max_orders[index] = order
+                for keys, order in self._max_order.items():
+                    kk = [keys.index(key) for key in keys]
+                    ok = True
+                    for k in range(len(self._orders)):
+                        if index[k] > 0 and k not in kk: ok = False
+                        if index[k] == 0 and k in kk: ok = False
+                    if ok:
+                        max_orders[index] = order
+        else:
+            max_orders = self._max_order
+
+        self._mask = sumorder <= max_orders
 
     def getProperties(self, star):
         return np.array([star.data[k] for k in self._keys], dtype=float)
@@ -558,7 +593,6 @@ def basis(self, star):
             p[1:] = vals[i]
             pows1d.append(np.cumprod(p))
         # Use trick to produce outer product of all these powers
-        #pows2d = np.prod(np.ix_(*pows1d))
         pows2d = np.prod(np.meshgrid(*pows1d, indexing='ij'), axis=0)
         # Return linear array of terms making total power constraint
         return pows2d[self._mask]
@@ -595,4 +629,3 @@ def _finish_read(self, reader):
         data = reader.read_table('solution')
         assert data is not None
         self.q = data['q'][0]
-

tests/test_pixel.py

@@ -1728,6 +1728,15 @@ def test_color():
     piff.piffify(config)
     psf = piff.read(psf_file)
 
+    # Show that the basis includes cross terms:
+    np.testing.assert_equal(psf.interp._orders, [2,2,1])
+    assert psf.interp._max_order == 2
+    s = psf.stars[0]
+    np.testing.assert_allclose(
+        psf.interp.basis(s),
+        [ 1, s['color'], s['v'], s['v']*s['color'], s['v']**2,
+            s['u'], s['u']*s['color'], s['u']*s['v'], s['u']**2 ])
+
     for s in psf.stars:
         orig_stamp = s.image
         weight = s.weight
@@ -1742,6 +1751,35 @@ def test_color():
                                   # Anyway, I think the fit is working, just this test doesn't
                                   # seem quite the right thing.
 
+    # Repeat without the cross-terms between color and u/v.
+    config['psf']['interp']['max_order'] = { ('u','v') : 2 }
+    piff.piffify(config)
+    psf = piff.read(psf_file)
+
+    # Show that the basis now doesn't include cross terms:
+    np.testing.assert_equal(psf.interp._orders, [2,2,1])
+    assert psf.interp._max_order ==  { ('u','v') : 2 }
+    s = psf.stars[0]
+    print(s['u'], s['v'], s['color'])
+    print('basis = ',psf.interp.basis(s))
+    np.testing.assert_allclose(
+        psf.interp.basis(s),
+        [ 1, s['color'], s['v'], s['v']**2, s['u'], s['u']*s['v'], s['u']**2 ])
+
+    # Still works just as well without the cross terms.
+    for s in psf.stars:
+        orig_stamp = s.image
+        weight = s.weight
+        offset = s.center_to_offset(s.fit.center)
+        image = psf.draw(x=s['x'], y=s['y'], color=s['color'],
+                         stamp_size=32, flux=s.fit.flux, offset=offset)
+        resid = image - orig_stamp
+        chisq = np.sum(resid.array**2 * weight.array)
+        dof = np.sum(weight.array != 0)
+        print('color = ',s['color'],'chisq = ',chisq,'dof = ',dof)
+        assert chisq < dof * 1.5
+
+
 @timer
 def test_convert_func():
     """Test PixelGrid fitting with a non-trivial convert_func