polynomial.py

#!/usr/bin/env python
##############################################################################
#
# SrMise            by Luke Granlund
#                   (c) 2014 trustees of the Michigan State University
#                   (c) 2024 trustees of Columbia University in the City of New York
#                   All rights reserved.
#
# File coded by:    Luke Granlund
#
# See LICENSE.txt for license information.
#
##############################################################################

import logging

import numpy as np

from diffpy.srmise.baselines.base import BaselineFunction
from diffpy.srmise.srmiseerrors import SrMiseEstimationError

logger = logging.getLogger("diffpy.srmise")


class Polynomial(BaselineFunction):
    """Methods for evaluation and parameter estimation of a polynomial baseline."""

    def __init__(self, degree, Cache=None):
        """Initialize a polynomial function of degree d.

        Parameters
        degree: The degree of the polynomial.  Any negative value is interpreted
                as the polynomial of negative infinite degree.
        Cache: A class (not instance) which implements caching of BaseFunction
               evaluations.
        """
        # Guarantee valid degree
        try:
            self.degree = int(str(degree))
        except ValueError:
            emsg = "Argument degree must be an integer."
            raise ValueError(emsg)
        if self.degree < 0:
            self.degree = -1  # interpreted as negative infinity
        # Define parameterdict
        # e.g. {"a_0":3, "a_1":2, "a_2":1, "a_3":0} if degree is 3.
        parameterdict = {}
        for d in range(self.degree + 1):
            parameterdict["a_" + str(d)] = self.degree - d
        formats = ["internal"]
        default_formats = {"default_input": "internal", "default_output": "internal"}
        metadict = {"degree": (degree, repr)}
        BaselineFunction.__init__(self, parameterdict, formats, default_formats, metadict, None, Cache)

    # Methods required by BaselineFunction ####

    def estimate_parameters(self, r, y):
        """Estimate parameters for polynomial baseline.

        Estimation is currently implemented only for degree < 2.  This
        very rudimentary method assumes the baseline crosses the origin, and
        y=baseline+signal, where signal is primarily positive.

        Parameters
        r: (Numpy array) Data along r from which to estimate
        y: (Numpy array) Data along y from which to estimate

        Returns Numpy array of parameters in the default internal format.
        Raises NotImplementedError if estimation is not implemented for this
        degree, or SrMiseEstimationError if parameters cannot be estimated for
        any other reason."""
        if self.degree > 1:
            emsg = "Polynomial implements estimation for baselines of degree <= 1 only."
            raise NotImplementedError(emsg)
        if len(r) != len(y):
            emsg = "Arrays r, y must have equal length."
            raise ValueError(emsg)

        if self.degree == -1:
            return np.array([])

        if self.degree == 0:
            return np.array([0.0])

        if self.degree == 1:
            # Estimate degree=1 baseline.
            # Find best slope for y=slope*r using only the least 10% of all
            # points, assuming the non-baseline component of the data largely
            # lies above the baseline.
            # TODO: Make this more sophisticated.
            try:
                cut = np.max([len(y) / 10, 1])
                cut_idx = y.argsort()[:cut]

                import numpy.linalg as la

                a = np.array([r[cut_idx]]).T
                slope = la.lstsq(a, y[cut_idx])[0][0]
                return np.array([slope, 0.0])
            except Exception as e:
                emsg = "Error during estimation -- " + str(e)
                raise SrMiseEstimationError(emsg)

    def _jacobianraw(self, pars, r, free):
        """Return the Jacobian of a polynomial.

        Parameters
        pars: Sequence of parameters for a polynomial of degree d
        pars[0] = a_degree
        pars[1] = a_(degree-1)
        ...
        pars[d] = a_0
        r: sequence or scalar over which pars is evaluated
        free: sequence of booleans which determines which derivatives are
              needed.  True for evaluation, False for no evaluation.
        """
        if len(pars) != self.npars:
            emsg = "Argument pars must have " + str(self.npars) + " elements."
            raise ValueError(emsg)
        if len(free) != self.npars:
            emsg = "Argument free must have " + str(self.npars) + " elements."
            raise ValueError(emsg)
        jacobian = [None for p in range(self.npars)]
        if (free is False).sum() == self.npars:
            return jacobian

        # The partial derivative with respect to the nth coefficient of a
        # polynomial is just x^nth.
        for idx in range(self.npars):
            if free[idx]:
                jacobian[idx] = np.power(r, idx)
        return jacobian

    def _transform_parametersraw(self, pars, in_format, out_format):
        """Convert parameter values from in_format to out_format.

        Parameters
        pars: Sequence of parameters
        in_format: A format defined for this class
        out_format: A format defined for this class

        Defined Formats
        internal: [a_degree, a_(degree-1), ..., a_0]"""
        temp = np.array(pars)

        # Convert to intermediate format "internal"
        if in_format == "internal":
            pass
        else:
            raise ValueError("Argument 'in_format' must be one of %s." % self.parformats)

        # Convert to specified output format from "internal" format.
        if out_format == "internal":
            pass
        else:
            raise ValueError("Argument 'out_format' must be one of %s." % self.parformats)
        return temp

    def _valueraw(self, pars, r):
        """Return value of polynomial for the given parameters and r values.

        Parameters
        pars: Sequence of parameters for a polynomial of degree d
              pars[0] = a_degree
              pars[1] = a_(degree-1)
              ...
              pars[d] = a_0
              If degree is negative infinity, pars is an empty sequence.
        r: sequence or scalar over which pars is evaluated"""
        if len(pars) != self.npars:
            emsg = "Argument pars must have " + str(self.npars) + " elements."
            raise ValueError(emsg)
        return np.polyval(pars, r)

    def getmodule(self):
        return __name__


# end of class Polynomial

# simple test code
if __name__ == "__main__":

    # Test polynomial of degree 3
    print("Testing degree 3 polynomial")
    print("---------------------------")
    f = Polynomial(degree=3)
    r = np.arange(5)
    pars = np.array([3, 0, 1, 2])
    free = np.array([True, False, True, True])
    val = f._valueraw(pars, r)
    jac = f._jacobianraw(pars, r, free)
    print("Value:\n", val)
    print("Jacobian: ")
    for j in jac:
        print(" %s" % j)

    # Test polynomial of degree -oo
    print("\nTesting degree -oo polynomial (== 0)")
    print("------------------------------------")
    f = Polynomial(degree=-1)
    r = np.arange(5)
    pars = np.array([])
    free = np.array([])
    val = f._valueraw(pars, r)
    jac = f._jacobianraw(pars, r, free)
    print("Value:\n", val)
    print("Jacobian: ")
    for j in jac:
        print(" %s" % j)

    # Test linear estimation
    print("\nTesting linear baseline estimation")
    print("------------------------------------")
    f = Polynomial(degree=1)
    pars = np.array([1, 0])
    r = np.arange(0, 10, 0.1)
    y = -r + 10 * np.exp(-((r - 5) ** 2)) + np.random.rand(len(r))
    est = f.estimate_parameters(r, y)
    print("Actual baseline: ", np.array([-1, 0.0]))
    print("Estimated baseline: ", est)