Merge pull request #40 from ToFuProject/Issue033_PolynomialBck

background `linear` => `poly` (deg = 2)

Author: dvezinet

spectrally/_class00_plot.py

@@ -267,7 +267,7 @@ def plot_axvlines(
     # if ne_scale is None:
         # ne_scale = 'log'
     # if Te_scale is None:
-        # Te_scale = 'linear'
+        # Te_scale = 'poly'
 
     # return
 
@@ -506,4 +506,4 @@ def plot_dominance_map(
         )
 
     return dax
-"""
+"""

spectrally/_class01_SpectralModel.py

@@ -61,15 +61,15 @@ def add_spectral_model(
                 'key': {'ref': k0, k1: [c0, c1], k2: [c0, c1]}
 
         Available function types for dmodel are:
-            - 'linear': typically a linear background
+            - 'poly': typically a polynomial background
             - 'exp': typically an exponential background
             - 'gauss': typically a doppler-broadened line
             - 'lorentz': ypically a natural-broadened line
             - 'pvoigt': typically a doppler-and-natural broadened line
 
         dconstraints holds a dict of constraints groups.
         Each group is a dict with a 'ref' variable
-        Other variables (keys) are compued as linear functions of 'ref'
+        Other variables (keys) are compued as poly functions of 'ref'
 
         Parameters
         ----------

spectrally/_class01_check_model.py

@@ -101,7 +101,7 @@ def _dmodel_err(key, dmodel):
             pstr = ", ".join([f"'{tpar[0]}': {tpar[1]}" for tpar in lpar])
             stri = f"\t- 'f{ii}': " + "{" + f"'type': '{k0}', {pstr}" + "}"
 
-        if k0 == 'linear':
+        if k0 == 'poly':
             lstr.append("\t# background-oriented")
         elif k0 == 'gauss':
             lstr.append("\t# spectral lines-oriented")
@@ -183,7 +183,7 @@ def _check_dmodel(
         # check key
 
         if isinstance(k0, int):
-            if typ in ['linear', 'exp']:
+            if typ in ['poly', 'exp']:
                 k1 = f'bck{ibck}'
                 ibck += 1
             else:

spectrally/_class01_display_models.py

@@ -140,7 +140,7 @@ def _get_xx(xx, ftype):
     # set if None
     # -------------
 
-    lbck = ['linear', 'exp_lamb']
+    lbck = ['poly', 'exp_lamb']
     llines = ['gauss', 'lorentz', 'pvoigt', 'voigt']
     lpulse = ['pulse_exp', 'pulse_gauss', 'lognorm']
 
@@ -201,12 +201,13 @@ def _get_dpar_xfree(ftype, xx):
     # xfree
     # ---------------
 
-    if ftype == 'linear':
+    if ftype == 'poly':
 
-        a1 = -1/Dx
-        a0 = 2 - a1 * x0
+        a0 = 1
+        a1 = -0.2
+        a2 = -1
 
-        xfree = np.r_[a0, a1]
+        xfree = np.r_[a0, a1, a2]
 
     elif ftype == 'exp_lamb':
 

spectrally/_class01_dparams.py

@@ -184,7 +184,7 @@ def _get_scale(
     # bck
     # ------------
 
-    if typ == 'linear':
+    if typ == 'poly':
 
         if var == 'c0':
             pass
@@ -614,4 +614,4 @@ def fit12d_dscales(dscales=None, dinput=None):
         indbs=dinput['valid']['indbs'],
     )
     return dscales
-"""
+"""

spectrally/_class01_fit_func_1d.py

@@ -49,6 +49,7 @@ def func(
         # for pulses
         lamb00 = lamb[0]
         lambD = lamb[-1] - lamb[0]
+        lambm = 0.5*(lamb[-1] + lamb[0])
 
         # iok
         if iok is not None:
@@ -75,16 +76,22 @@ def func(
             x_full = c2.dot(x_free**2) + c1.dot(x_free) + c0
 
         # ------------------
-        # sum all linear
+        # sum all poly
 
-        kfunc = 'linear'
+        kfunc = 'poly'
         if dind.get(kfunc) is not None:
 
             a0 = x_full[dind[kfunc]['a0']['ind']][:, None]
             a1 = x_full[dind[kfunc]['a1']['ind']][:, None]
+            a2 = x_full[dind[kfunc]['a2']['ind']][:, None]
 
             ind = dind['func'][kfunc]['ind']
-            val[ind, ...] = a0 + lamb * a1
+            lamb_rel = (lamb - lambm) / lambD
+            val[ind, ...] = (
+                a0
+                + a1 * lamb_rel
+                + a2 * lamb_rel**2
+            )
 
         # --------------------
         # sum all exponentials
@@ -426,6 +433,7 @@ def func(
         # for pulses
         lamb00 = lamb[0]
         lambD = lamb[-1] - lamb[0]
+        lambm = 0.5*(lamb[-1] + lamb[0])
 
         # iok
         if iok is not None:
@@ -453,11 +461,13 @@ def func(
             x_full = c2.dot(x_free**2) + c1.dot(x_free) + c0
 
         # -------
-        # linear
+        # poly
 
-        kfunc = 'linear'
+        kfunc = 'poly'
         if dind.get(kfunc) is not None:
 
+            lamb_rel = (lamb - lambm) / lambD
+
             vind = dind['jac'][kfunc].get('a0')
             if vind is not None:
                 ival, ivar = vind['val'], vind['var']
@@ -466,7 +476,12 @@ def func(
             vind = dind['jac'][kfunc].get('a1')
             if vind is not None:
                 ival, ivar = vind['val'], vind['var']
-                val[:, ival] = lamb * scales[None, ival]
+                val[:, ival] = lamb_rel * scales[None, ival]
+
+            vind = dind['jac'][kfunc].get('a2')
+            if vind is not None:
+                ival, ivar = vind['val'], vind['var']
+                val[:, ival] = lamb_rel**2 * scales[None, ival]
 
         # --------
         # exp_lamb

spectrally/_class01_model_dict.py

@@ -17,7 +17,7 @@
 
 _LMODEL_ORDER = [
     # background
-    'linear', 'exp_lamb', # 'exp_E',
+    'poly', 'exp_lamb', # 'exp_E',
     # spectral lines
     'gauss', 'lorentz', 'pvoigt', 'voigt',
     # pulse shape
@@ -38,13 +38,15 @@
     # -------------------
 
     # -----------
-    # linear
+    # poly
 
-    'linear': {
-        'var': ['a0', 'a1'],
-        'description': 'straight line',
+    'poly': {
+        'var': ['a0', 'a1', 'a2'],
+        'description': 'polynomial (up to deg = 2)',
         'expressions': {
-            'main': r"$a_0 + a_1\lambda$",
+            'main': r"$\left\{ \begin{array}{ll} a_0 + a_1\frac{\lambda - <\lambda>}{\Delta\lambda} + a_2\left(\frac{\lambda - <\lambda>}{\Delta\lambda}\right)^2\\ <\lambda> = \frac{\lambda[0] + \lambda[-1]}{2} \\ \Delta\lambda = \lambda[-1] - \lambda[0]\end{array} \right.$",
+            'argmax': r"$\lambda_{max} = <\lambda> - \frac{a_1}{2a_2}\Delta\lambda$",
+            'max': r"$f(\lambda_{max}) = a_0 - \frac{a_1^2}{4a_2}$",
         },
     },
 

spectrally/_class01_moments.py

@@ -216,6 +216,7 @@ def func(
         # prepare
 
         lambD = lamb[-1] - lamb[0]
+        lambm = 0.5*(lamb[0] + lamb[-1])
 
         # ----------
         # initialize
@@ -268,13 +269,14 @@ def func(
                     dout[kfunc][kvar] = extract(kfunc, kvar)
 
         # ------------------
-        # sum all linear
+        # sum all poly
 
-        kfunc = 'linear'
+        kfunc = 'poly'
         if dind.get(kfunc) is not None:
 
             a0 = dout[kfunc]['a0']
             a1 = dout[kfunc]['a1']
+            a2 = dout[kfunc]['a2']
 
             # integral
             if lamb is not None:
@@ -283,6 +285,13 @@ def func(
                     + a1 * (lamb[-1]**2 - lamb[0]**2)/2
                 )
 
+            # argmax, max
+            dout[kfunc]['argmax'] = np.full(a0.shape, np.nan)
+            dout[kfunc]['max'] = np.full(a1.shape, np.nan)
+            iok = a2 != 0
+            dout[kfunc]['argmax'][iok] = lambm - lambD * a1[iok]/(2*a2[iok])
+            dout[kfunc]['max'][iok] = a0[iok] - a1[iok]**2 / (4*a2[iok])
+
         # --------------------
         # sum all exponentials
 

spectrally/_class01_show.py

@@ -239,7 +239,7 @@ def get_available_spectral_model_functions(
     for kf in lkeys:
 
         # use case
-        if kf == 'linear':
+        if kf == 'poly':
             use = 'bck'
         elif kf == 'gauss':
             use = 'lines'

spectrally/_class02_x0_scale_bounds.py

@@ -177,22 +177,39 @@ def _get_scales_bounds(
     ldins = [(dscales, scales), (dbounds_low, bounds0), (dbounds_up, bounds1)]
 
     # ------------------
-    # all linear
+    # all poly
     # ------------------
 
-    kfunc = 'linear'
+    kfunc = 'poly'
     if dind.get(kfunc) is not None:
 
         a1max = (data_max - data_min) / lambD
 
+        # -------
+        # a0
+
+        kvar = 'a0'
+        vind = dind['jac'][kfunc].get(kvar)
+        if vind is not None:
+            ival, ivar = vind['val'], vind['var']
+            scales[ival] = max(np.abs(data_max), np.abs(data_min))
+            bounds0[ival] = -10.
+            bounds1[ival] = 10.
+
+            for ii, (din, val) in enumerate(ldins):
+                _update_din_from_user(
+                    din, kfunc, kvar, val,
+                    scales=None if ii == 0 else scales
+                )
+
         # -------
         # a1
 
         kvar = 'a1'
         vind = dind['jac'][kfunc].get(kvar)
         if vind is not None:
             ival, ivar = vind['val'], vind['var']
-            scales[ival] = a1max
+            scales[ival] = max(np.abs(data_max), np.abs(data_min))
             bounds0[ival] = -10.
             bounds1[ival] = 10.
 
@@ -203,16 +220,13 @@ def _get_scales_bounds(
                 )
 
         # -------
-        # a0
+        # a2
 
-        kvar = 'a0'
+        kvar = 'a2'
         vind = dind['jac'][kfunc].get(kvar)
         if vind is not None:
             ival, ivar = vind['val'], vind['var']
-            scales[ival] = max(
-                np.abs(data_min - a1max*lamb[0]),
-                np.abs(data_max + a1max*lamb[0]),
-            )
+            scales[ival] = max(np.abs(data_max), np.abs(data_min))
             bounds0[ival] = -10.
             bounds1[ival] = 10.
 
@@ -703,13 +717,28 @@ def _get_x0(
     ldins = [(dx0, x0)]
 
     # ------------------
-    # all linear
+    # all poly
     # ------------------
 
-    kfunc = 'linear'
+    kfunc = 'poly'
     if dind.get(kfunc) is not None:
 
-        a1max = ((data_max - data_min) / lambD) / 10
+        # a1max = ((data_max - data_min) / lambD) / 10
+
+        # -------
+        # a0
+
+        kvar = 'a0'
+        vind = dind['jac'][kfunc].get(kvar)
+        if vind is not None:
+            ival, ivar = vind['val'], vind['var']
+            x0[ival] = data_mean / scales[ival]
+
+            for ii, (din, val) in enumerate(ldins):
+                _update_din_from_user(
+                    din, kfunc, kvar, val,
+                    scales=scales,
+                )
 
         # -------
         # a1
@@ -718,7 +747,7 @@ def _get_x0(
         vind = dind['jac'][kfunc].get(kvar)
         if vind is not None:
             ival, ivar = vind['val'], vind['var']
-            x0[ival] = a1max / scales[ival]
+            x0[ival] = 0. / scales[ival]
 
             for ii, (din, val) in enumerate(ldins):
                 _update_din_from_user(
@@ -727,13 +756,13 @@ def _get_x0(
                 )
 
         # -------
-        # a0
+        # a2
 
-        kvar = 'a0'
+        kvar = 'a2'
         vind = dind['jac'][kfunc].get(kvar)
         if vind is not None:
             ival, ivar = vind['val'], vind['var']
-            x0[ival] = np.abs(data_min - a1max*lamb[0]) / scales[ival]
+            x0[ival] = 0. / scales[ival]
 
             for ii, (din, val) in enumerate(ldins):
                 _update_din_from_user(

spectrally/tests/_hxr_pulses.py

@@ -97,7 +97,7 @@ def main(
     coll.add_spectral_model(
         key='sm_exp',
         dmodel={
-            'bck': 'linear',
+            'bck': 'poly',
             'pulse': 'pulse_exp',
         },
         dconstraints={
@@ -109,7 +109,7 @@ def main(
     coll.add_spectral_model(
         key='sm_gauss',
         dmodel={
-            'bck': 'linear',
+            'bck': 'poly',
             'pulse': 'pulse_gauss',
         },
         dconstraints={
@@ -121,7 +121,7 @@ def main(
     coll.add_spectral_model(
         key='sm_lognorm',
         dmodel={
-            'bck': 'linear',
+            'bck': 'poly',
             'pulse': 'lognorm',
         },
         dconstraints={