Ticket #312: slit_test.py

r"""
Explore different integration schemes for slit smearing.

Introduction
============

Slit smearing evaluates the following integral:

.. math::

    I_s(q_o) = \int_0^{\Delta q_v} I(\sqrt{q_o^2 + u^2}) du

where $\Delta q_v$ is the vertical slit contribution. See
`http://www.ncnr.nist.gov/staff/hammouda/distance_learning/chapter_61.pdf`_

Procedure
=========

This program is not intended as an end user application, but rather as a
platform for exploring implementations of the resolution integral.
Different experiments can be performed by modifying the :func:`demo` function.
The integration functions have the name smear_METHOD, for a variety of
different methods.  Each one takes a form factor with parameters, a set of
q values, a resolution, and maybe a set of q_calc values at which to evaluate
the theory function.   The *ms* and *mt* variables in :func:`demo`, determine
which methods are being compared.  :func:`interpolate` allows you to fill
in the integral between q points since the demo function is rapidly
oscillating.  *TEST_DATA_SLIT_SMEAR* contains the values calculated in
Igor for a sphere model with a radius of 6 microns and a contrast of 3e-6
inverse Angstroms.

:func:`smear_romberg` is a direct implementation of the integral using
adaptive romberg integration from scipy.  This method would add considerable
complexity to sasmodels since we could no longer rely on evaluating the
models at a fixed set of q points for all parameter sets in a fit.  A
"good enough" method with fixed q, although less accurate and requiring
more q points on average to evaluate, will be a win in complexity and perhaps
performance.  The romberg integration parameters are set to produce a
pretty good value for the integral.

The remaining methods used a fixed set of *q* points, *q_calc*, to evaluate the
integral.   They are :func:`smear_simpsons`,  :func:`smear_trapezoid`,
and :func:`smear_rectangular`.  Except for :func:`smear_simpsons_pow`,
*q_calc* is extrapolated to the full width of the resolution function.
The extrapolation uses exponential spacing between the points, much like *q*
spacing within the measured data range, though :func:`smear_simpsons_lin`
uses linear spacing.

:func:`smear_rectangular` is implemented as a non-square weight matrix which
can be multiplied by a set of computed values *I(q_calc)* to produce
*I_smeared(q)*.  The same style could be used for trapezoid and Simpson's
rule with the appropriate bookkeeping.  This was not done because at least
for the demo problem, rectangular performed better than the other
non-adaptive methods, and so wins on both simplicity and accuracy.

:func:`smear_simpsons_pow` uses simpson's rule to estimate the integral within
the range of the data, and a power law approximation for the tail above the
maximum q value.  This is akin to the Igor implementation of slit smearing,
though it uses the theory function itself to estimate the power law rather
than data.

The power law approximation has two problems:

1. the power is estimated from the data, which means that this will only
   work if the data has a flat region at the tail, and
2. any features such as lamellar peaks that occur beyond the end of the
   data will not be correctly modeled.

The demo shows that this method performs poorly for high q when oscillations
in the theory have not completely decayed by the end of the q range, which
is not really a surprise.  It is not a problem in practice for usans data on
igor since polydispersity will likely dampen any ringing that a monodisperse
model might exhibit.  The demo does not simulate a model containing a peak at
q beyond the end of the data.


Conclusions
===========

At least for this theory function, the simple rectangular integral is as
good as the more complicated trapezoid and Simpson's rule integration.
Due to the oscillations in the theory function, the relative error is
huge (50%) at high q.  Oversampling by a factor of 10 (requiring 1020
evaluations rather than 122), reduces the error to 2.5%.  Another factor
of 10 (9130 evaluations) reduces error to 0.15%. Getting error down to
5e-5 requires about 90,000 evaluations.

In comparison, romberg integration requires only 1360 evaluations.

"""
import numpy as np
from numpy import sqrt, cos, sin, pi, log10, log, exp

def smear_romberg(q, form, pars, delta_qv, q_calc=None):
    """
    Romberg integration.

    This is an adaptive integration technique.  It is called with settings
    that make it slow to evaluate but give it good accuracy.
    """
    from scipy.integrate import romberg
    evaluations = [0]
    def _fn(u, q0):
        evaluations[0] += 1 if np.isscalar(q0) else len(q0)
        return form(sqrt(q0**2 + u**2), *pars)
    r = [romberg(_fn, 0, delta_qv, args=(qi,),
                 divmax=100, vec_func=True, tol=0, rtol=1e-14)
         for qi in q]
    print "romberg evaluations:", evaluations[0]
    return r


def smear_simpsons_pow(q, form, pars, delta_qv, q_calc=None):
    """
    Use Simpson's rule with power low approximation to compute the integral.

    The *q* range between the end of the data and *delta_qv* is not sampled.
    Instead it is approximated by a power law computed from the (I(delta_qv)) and
    the final 25\% of the theory function.  Igor fits the data to the a power
    law and assumes that it will be an adequate estimate for the fitted
    theory function.
    """
    from scipy.integrate import simps
    if q_calc is None: q_calc = q
    x = np.hstack((q_calc, q_calc[-1]+delta_qv))
    print "#q:",q.size, "#q_calc:",x.size
    y = form(x, *pars)
    A, m = fit_power_law(x, y, portion=0.25)
    return [simps(y[i:-1], q_to_u(x[i:-1], q[i]))
            + power_law_residual(delta_qv, q[i], q[-1], A, m)
            for i in range(len(q))]


def smear_simpsons_lin(q, form, pars, delta_qv, q_calc=None):
    """
    Use Simpson's rule with linear extrapolation to compute the integral.

    The evaluation points *q* are extrapolated to *delta_qv+q[-1]* using
    the final *q* step as the step size.

    For a typical NCNR USANS dataset this requires over 10000 extra evaluations
    of the function.
    """
    from scipy.integrate import simps
    if q_calc is None: q_calc = q
    x = linear_extrapolation(q_calc, q_calc[-1]+delta_qv)
    print "#q:",q.size, "#q_calc:",x.size
    y = form(x, *pars)
    return [simps(y[j:], q_to_u(x[j:], qi))
            for qi in q for j in [np.searchsorted(x,qi)]]


def smear_simpsons(q, form, pars, delta_qv, q_calc=None):
    """
    Use Simpson's rule with geometric extrapolation to compute the integral.

    The evaluation points *q* are extrapolated to *delta_qv+q[-1]* using
    the geometric average of *q* as the step size.

    For a typical NCNR USANS dataset the function is evaluated at twice as
    many points as there are points in q.
    """
    from scipy.integrate import simps
    if q_calc is None: q_calc = q
    x = geometric_extrapolation(q_calc, q_calc[-1]+delta_qv)
    print "#q:",q.size, "#q_calc:",x.size
    y = form(x, *pars)
    return [simps(y[j:], q_to_u(x[j:], qi))
            for qi in q for j in [np.searchsorted(x,qi)]]


def smear_trapezoid(q, form, pars, delta_qv, q_calc=None):
    """
    Use the trapezoid rule with geometric extrapolation to compute the integral.
    """
    from scipy.integrate import trapz

    if q_calc is None: q_calc = q
    x = geometric_extrapolation(q_calc, q_calc[-1]+delta_qv)
    print "#q:",q.size, "#q_calc:",x.size
    y = form(x, *pars)
    return [trapz(y[j:], q_to_u(x[j:], qi))
            for qi in q for j in [np.searchsorted(x,qi)]]


def smear_rectangular(q, form, pars, delta_qv, q_calc=None):
    """
    Use the trapezoid rule with geometric extrapolation to compute the integral.
    """
    from scipy.integrate import trapz
    if q_calc is None: q_calc = q
    x = geometric_extrapolation(q_calc, q_calc[-1]+delta_qv)
    print "#q:",q.size, "#q_calc:",x.size
    edges = bin_edges(x)
    W = np.zeros((len(q), len(x)), 'd')
    for i, qi in enumerate(q):
        W[i,:] = np.diff(q_to_u(edges, qi))

    y = form(x, *pars)
    return np.dot(W,y)


###########################################################################
# Helper functions

def sphere(q, radius, contrast):
    """
    Sphere form factor with no polydispersity
    """
    qr = q * radius
    sn, cn = sin(qr), cos(qr)
    bes = 3 * (sn - qr * cn) / qr ** 3 # may be 0/0 but we fix that next line
    #bes[qr == 0] = 1
    fq = bes * contrast
    return 4*pi*radius**3/3 * 1.0e-4 * fq ** 2



def q_to_u(q, q0):
    """
    Convert *q* values to *u* values for the integral computed at *q0*.
    """
    # array([value])**2 - value**2 is not always zero
    qpsq = q**2 - q0**2
    qpsq[qpsq<0] = 0
    return sqrt(qpsq)


def interpolate(q, max_step):
    """
    Returns *q_calc* with points spaced at most max_step apart.
    """
    step = np.diff(q)
    index = step>max_step
    if np.any(index):
        inserts = []
        for q_i,step_i in zip(q[:-1][index],step[index]):
            n = np.ceil(step_i/max_step)
            inserts.extend(q_i + np.arange(1,n)*(step_i/n))
        # Extend a couple of fringes beyond the end of the data
        inserts.extend(q[-1] + np.arange(1,16)*max_step)
        q_calc = np.sort(np.hstack((q,inserts)))
    else:
        q_calc = q
    return q_calc

def linear_extrapolation(q, qmax):
    """
    Extrapolate *q* out to *qmax* using the final step in *q* as the step
    size.
    """
    extrapolate = np.arange(q[-1], qmax, q[-1]-q[-2])
    x = np.concatenate((q, extrapolate[1:]))


def geometric_extrapolation(q, qmax):
    r"""
    Extrapolate *q* out to *qmax* using geometric steps, with the average
    geometric step size in *q* as the step size.

    Starting at $q_1$ and stepping geometrically by $\Delta q$
    to $q_n$ in $n$ points gives a geometric average of:

    .. math::

         \log \Delta q = (\log q_n - log q_1) / (n - 1)

    From this we can compute the number of steps required to extend $q$
    from $q_n$ to $q_\text{max}$ by $\Delta q$ as:

    .. math::

         n_\text{extend} = (\log q_\text{max} - \log q_n) / \log \Delta q

    Substituting:

        n_\text{extend} = (n-1) (\log q_\text{max} - \log q_n)
            / (\log q_n - log q_1)
    """
    n_extend = (len(q)-1)*(log(qmax) - log(q[-1]))/(log(q[-1]) - log(q[0]))
    extrapolate = np.logspace(log10(q[-1]), log10(qmax), np.ceil(n_extend)+1)
    x = np.concatenate((q, extrapolate[1:]))
    return x


def fit_power_law(x, y, portion=1.0):
    """
    Fit a power law $y=Ax^{-m}$ against *x*, *y*.  If *portion* is given, use
    only the given portion from the tail of the data in the fit.
    """
    fitting_points = max(2, np.ceil(len(x)*portion))
    p = np.polyfit(log(x[-fitting_points:]), log(y[-fitting_points:]), 1)
    A,m = exp(p[1]), -p[0]
    #print "A,m",A,m


def power_law_residual(delta_qv, q0, qmax, A, m):
    r"""
    Residual term for the integration, assuming I(q) follows a power law for
    high q.

    This integrates $I_p(\sqrt(q_o^2 + u^2})$ by $u$ for $u$ corresponding to
    $q_\text{max}$ to $u=\Delta q_v$, where $I_p$ is:

    .. math::

        I_p(q) = A q^{-m}
    """
    from scipy.integrate import romberg
    a, b = sqrt(qmax**2 - q0**2), delta_qv
    _residual_fn = lambda u, q0, A, m: A*(q0**2 + u**2)**(-m/2)
    r = romberg(_residual_fn, a, b, args=(q0, A, m),
                divmax=100, vec_func=True, tol=0, rtol=1e-8)
    return r

def bin_edges(x):
    """
    Determine bin edges from bin centers, assuming that edges are centered
    between the bins.

    Note: this uses the arithmetic mean, which may not be appropriate for
    log-scaled data.
    """
    if len(x) < 2 or (np.diff(x)<0).any():
        raise ValueError("Expected bins to be an increasing set")
    edges = np.hstack([
        x[0]  - 0.5*(x[1]  - x[0]),  # first point minus half first interval
        0.5*(x[1:] + x[:-1]),        # mid points of all central intervals
        x[-1] + 0.5*(x[-1] - x[-2]), # last point plus half last interval
        ])
    return edges

def Iq_smeared(q, form, pars, scale=1, background=0., delta_qv=0.117,
               method='trapezoid', max_step=None):
    """
    Compute slit-smeared version of *form(q;pars)*
    """
    q_calc = interpolate(q, max_step) if max_step is not None else q
    integrator = globals()['smear_'+method]
    res = integrator(q, form, pars, delta_qv, q_calc)
    return scale/delta_qv*np.asarray(res) + background


def Iq(q, form, pars, scale=1, background=0):
    """
    Compute unsmeared version of *form(q;pars)*
    """
    return scale*form(q, *pars) + background


TEST_DATA_SLIT_SPHERE = """\
2.26097e-05 0.117 5.5781372896e+09 1.4626077708e+06
2.53847e-05 0.117 5.0363141458e+09 1.3117318023e+06
2.81597e-05 0.117 4.4850108103e+09 1.1594863713e+06
3.09347e-05 0.117 3.9364658459e+09 1.0094881630e+06
3.37097e-05 0.117 3.4019975074e+09 8.6518941303e+05
3.92597e-05 0.117 2.4139519814e+09 6.0232158311e+05
4.48097e-05 0.117 1.5816877820e+09 3.8739994090e+05
5.03597e-05 0.117 9.3715407224e+08 2.2745304775e+05
5.59097e-05 0.117 4.8387917428e+08 1.2101295768e+05
6.14597e-05 0.117 2.0193586928e+08 6.0055107771e+04
6.70097e-05 0.117 5.5886110911e+07 3.2749521065e+04
7.25597e-05 0.117 3.7782348010e+06 2.6350963616e+04
7.81097e-05 0.117 5.3407817904e+06 2.9624963314e+04
8.36597e-05 0.117 2.7975485523e+07 3.4403962254e+04
8.92097e-05 0.117 4.9845448282e+07 3.6130017903e+04
9.47597e-05 0.117 6.0092588905e+07 3.3495107849e+04
1.00310e-04 0.117 5.6823430831e+07 2.7475726279e+04
1.05860e-04 0.117 4.3857024036e+07 2.0144282226e+04
1.11410e-04 0.117 2.7277144760e+07 1.3647403260e+04
1.22510e-04 0.117 3.3119334113e+06 6.6519711526e+03
1.33610e-04 0.117 1.4412859402e+06 6.9726212813e+03
1.44710e-04 0.117 8.5620162463e+06 8.1441335775e+03
1.55810e-04 0.117 9.6957429033e+06 6.4559996521e+03
1.66910e-04 0.117 4.3818341914e+06 3.6252154156e+03
1.78010e-04 0.117 2.7448997387e+05 2.4006505342e+03
1.89110e-04 0.117 8.0472009936e+05 2.8187789089e+03
2.00210e-04 0.117 2.8149552834e+06 3.0915662855e+03
2.11310e-04 0.117 2.7510907861e+06 2.3722530293e+03
2.22410e-04 0.117 1.0053133293e+06 1.4473468311e+03
2.33510e-04 0.117 5.8428305052e+03 1.2048540556e+03
2.44610e-04 0.117 5.1699305004e+05 1.4625670042e+03
2.55710e-04 0.117 1.2120227268e+06 1.5010705973e+03
2.66810e-04 0.117 9.7896842846e+05 1.1336343426e+03
2.77910e-04 0.117 2.5507264791e+05 8.1848818080e+02
3.05660e-04 0.117 5.2403101181e+05 7.4913374357e+02
3.33410e-04 0.117 5.8699343809e+04 4.4669964560e+02
3.61160e-04 0.117 3.0844327150e+05 4.6774007542e+02
3.88910e-04 0.117 8.3360142970e+03 2.7169550220e+02
4.16660e-04 0.117 1.8630080583e+05 3.0710983679e+02
4.44410e-04 0.117 3.1616804732e-01 1.7959006831e+02
4.72160e-04 0.117 1.1299016314e+05 2.0763952339e+02
4.99910e-04 0.117 2.9952522747e+03 1.2536542765e+02
5.27660e-04 0.117 6.7625695649e+04 1.4013969777e+02
5.55410e-04 0.117 7.6927460089e+03 8.2145593180e+01
6.10910e-04 0.117 1.1229057779e+04 8.4519745643e+01
6.66410e-04 0.117 1.3035567943e+04 8.1554625609e+01
7.21910e-04 0.117 1.3309931343e+04 7.4437319172e+01
7.77410e-04 0.117 1.2462626212e+04 6.4697088261e+01
8.32910e-04 0.117 1.0912927143e+04 5.3773301044e+01
8.88410e-04 0.117 9.0172597469e+03 4.2843375753e+01
9.43910e-04 0.117 7.0496495917e+03 3.2771032724e+01
9.99410e-04 0.117 5.2030483682e+03 2.4113557144e+01
1.05491e-03 0.117 3.5988976711e+03 1.7160773658e+01
1.11041e-03 0.117 2.2996060652e+03 1.2016626459e+01
1.22141e-03 0.117 6.4766590598e+02 6.0373017740e+00
1.33241e-03 0.117 4.1963483264e+01 4.5215452974e+00
1.44341e-03 0.117 6.3370708246e+01 5.1054681903e+00
1.55441e-03 0.117 3.0736750577e+02 5.9176165298e+00
1.66541e-03 0.117 5.0327682399e+02 5.9815000189e+00
1.77641e-03 0.117 5.4084331454e+02 5.1634639625e+00
1.88741e-03 0.117 4.3488671756e+02 3.8535158148e+00
1.99841e-03 0.117 2.6322287860e+02 2.5824997753e+00
2.10941e-03 0.117 1.0793633150e+02 1.7315517194e+00
2.22041e-03 0.117 1.8474448850e+01 1.4077213604e+00
2.33141e-03 0.117 1.5864062279e+00 1.4771560682e+00
2.44241e-03 0.117 3.2267213848e+01 1.6916253448e+00
2.55341e-03 0.117 7.4289116207e+01 1.8274751193e+00
2.66441e-03 0.117 9.9000521929e+01 1.7706812289e+00
"""

def demo():
    from matplotlib import pyplot as plt

    ## ==== Set data and model parameters
    delta_qv = 0.117
    q = np.logspace(-5,-3,400)
    #q = np.logspace(log10(2.3e-5),log10(2.7e-3),68*10)
    #q = np.logspace(log10(2.3e-5),log10(2.7e-3),68)
    if 1:  # use values from Igor
        data = np.loadtxt(TEST_DATA_SLIT_SPHERE.split('\n')).T
        q, dq, _, answer = data
        delta_qv = dq[0]
    else:
        answer = None
    scale, background = 0.01, 0. #0.01
    pars = (60000, 3) # radius, contrast

    ## Set the maximum step size between q values for numerical integration
    ## This can be based on the biggest dimension in the model, the visible
    ## oscillation in the data, or None if the data is assumed to be j
    d_max = pars[0]
    points_per_fringe = 2000
    max_step = 2*pi/(d_max*points_per_fringe)
    #max_step = None
    #print "max_step",max_step

    ## ==== Select which models are being compared
    ## Models are smear_MODEL functions above.
    ## romberg is the slow accurate model computed with adaptive integration
    ## rectangular is the simplest model to implement, and hopefully good
    ## enough.  It is the only one that uses max_step at the time of this
    ## writing, and is what we will use in sasmodels since it seems to be
    ## good enough.
    #ms,mt = "simpsons","rectangular"
    ms,mt = "romberg","rectangular"
    #ms,mt = "romberg","trapezoid"
    #ms,mt = "romberg","simpsons"

    ## ==== Evaluate the model and its comparison
    q_calc = interpolate(q, max_step) if max_step is not None else q
    y = Iq(q_calc, sphere, pars, scale=scale, background=background)
    ys = Iq_smeared(q, sphere, pars, delta_qv=delta_qv, max_step=max_step,
                    scale=scale, background=background, method=ms)
    yt = Iq_smeared(q, sphere, pars, delta_qv=delta_qv, max_step=max_step,
                    scale=scale, background=background, method=mt)

    ## ==== plots
    plt.subplot(211)
    plt.loglog(q_calc, y, '-', label="unsmeared")
    plt.loglog(q, ys, '.', label=ms, hold=True)
    plt.loglog(q, yt, '.', label=mt, hold=True)
    if answer is not None:
        plt.loglog(q, answer, '-', label='igor', hold=True)
    plt.ylabel("y")
    ## Uncomment to see linear fringing in the model; this works particularly
    ## well when q_calc is set to extrapolated and interpolated points.
    #plt.xscale('linear')
    plt.grid(True)
    plt.legend()
    plt.subplot(212)
    err = (ys-yt)/ys
    print mt, "err max:", max(abs(err)), " median:", np.median(abs(err))
    plt.semilogx(q, err, '-', label=mt)
    ## Uncomment to compare test model to igor
    #yt=answer # use igor output for comparison
    #plt.semilogx(q, (ys-yt)/ys, '-', label='Igor', hold=True)
    plt.xlabel("q")
    plt.ylabel("(%s - %s)/%s"%(ms,mt,ms))
    plt.legend()


if __name__ == "__main__":
    demo()
    import pylab; pylab.show()