# Licensed under a 3-clause BSD style license - see LICENSE.rst

"""General purpose timer related functions."""

from __future__ import (absolute_import, division, print_function,

unicode_literals)

from ..extern import six

from ..extern.six.moves import range

# STDLIB

import time

import warnings

from collections import Iterable, OrderedDict

from functools import partial, wraps

# THIRD-PARTY

import numpy as np

# LOCAL

from .. import units as u

from .. import log

from .. import modeling

from .exceptions import AstropyUserWarning

__all__ = ['timefunc', 'RunTimePredictor']

__doctest_skip__ = ['timefunc']

def timefunc(num_tries=1, verbose=True):

"""Decorator to time a function or method.

Parameters

----------

num_tries : int, optional

Number of calls to make. Timer will take the

average run time.

verbose : bool, optional

Extra log information.

Returns

-------

tt : float

Average run time in seconds.

result

Output(s) from the function.

Examples

--------

To add timer to time `numpy.log` for 100 times with

verbose output::

import numpy as np

from astropy.utils.timer import timefunc

@timefunc(100)

def timed_log(x):

return np.log(x)

To run the decorated function above:

>>> t, y = timed_log(100)

INFO: timed_log took 9.29832458496e-06 s on AVERAGE for 100 call(s). [...]

>>> t

9.298324584960938e-06

>>> y

4.6051701859880918

"""

def real_decorator(function):

@wraps(function)

def wrapper(*args, **kwargs):

ts = time.time()

for i in range(num_tries):

result = function(*args, **kwargs)

te = time.time()

tt = (te - ts) / num_tries

if verbose: # pragma: no cover

log.info('{0} took {1} s on AVERAGE for {2} call(s).'.format(

function.__name__, tt, num_tries))

return tt, result

return wrapper

return real_decorator

class RunTimePredictor(object):

"""Class to predict run time.

.. note:: Only predict for single varying numeric input parameter.

Parameters

----------

func : function

Function to time.

args : tuple

Fixed positional argument(s) for the function.

kwargs : dict

Fixed keyword argument(s) for the function.

Examples

--------

>>> from astropy.utils.timer import RunTimePredictor

Set up a predictor for :math:`10^{x}`:

>>> p = RunTimePredictor(pow, 10)

Give it baseline data to use for prediction and

get the function output values:

>>> p.time_func(range(10, 1000, 200))

>>> for input, result in sorted(p.results.items()):

... print("pow(10, {0})\\n{1}".format(input, result))

pow(10, 10)

10000000000

pow(10, 210)

10000000000...

pow(10, 410)

10000000000...

pow(10, 610)

10000000000...

pow(10, 810)

10000000000...

Fit a straight line assuming :math:`\\text{arg}^{1}` relationship

(coefficients are returned):

>>> p.do_fit() # doctest: +SKIP

array([1.16777420e-05, 1.00135803e-08])

Predict run time for :math:`10^{5000}`:

>>> p.predict_time(5000) # doctest: +SKIP

6.174564361572262e-05

Plot the prediction:

>>> p.plot(xlabeltext='Power of 10') # doctest: +SKIP

.. image:: /_static/timer_prediction_pow10.png

:width: 450px

:alt: Example plot from `astropy.utils.timer.RunTimePredictor`

When the changing argument is not the last, e.g.,

:math:`x^{2}`, something like this might work:

>>> p = RunTimePredictor(lambda x: pow(x, 2))

>>> p.time_func([2, 3, 5])

>>> sorted(p.results.items())

[(2, 4), (3, 9), (5, 25)]

"""

def __init__(self, func, *args, **kwargs):

self._funcname = func.__name__

self._pfunc = partial(func, *args, **kwargs)

self._cache_good = OrderedDict()

self._cache_bad = []

self._cache_est = OrderedDict()

self._cache_out = OrderedDict()

self._fit_func = None

self._power = None

@property

def results(self):

"""Function outputs from `time_func`.

A dictionary mapping input arguments (fixed arguments

are not included) to their respective output values.

"""

return self._cache_out

@timefunc(num_tries=1, verbose=False)

def _timed_pfunc(self, arg):

"""Run partial func once for single arg and time it."""

return self._pfunc(arg)

def _cache_time(self, arg):

"""Cache timing results without repetition."""

if arg not in self._cache_good and arg not in self._cache_bad:

try:

result = self._timed_pfunc(arg)

except Exception as e:

warnings.warn(str(e), AstropyUserWarning)

self._cache_bad.append(arg)

else:

self._cache_good[arg] = result[0] # Run time

self._cache_out[arg] = result[1] # Function output

def time_func(self, arglist):

"""Time the partial function for a list of single args

and store run time in a cache. This forms a baseline for

the prediction.

This also stores function outputs in `results`.

Parameters

----------

arglist : list of numbers

List of input arguments to time.

"""

if not isinstance(arglist, Iterable):

arglist = [arglist]

# Preserve arglist order

for arg in arglist:

self._cache_time(arg)

# FUTURE: Implement N^x * O(log(N)) fancy fitting.

def do_fit(self, model=None, fitter=None, power=1, min_datapoints=3):

"""Fit a function to the lists of arguments and

their respective run time in the cache.

By default, this does a linear least-square fitting

to a straight line on run time w.r.t. argument values

raised to the given power, and returns the optimal

intercept and slope.

Parameters

----------

model : `astropy.modeling.Model`

Model for the expected trend of run time (Y-axis)

w.r.t. :math:`\\text{arg}^{\\text{power}}` (X-axis).

If `None`, will use `~astropy.modeling.polynomial.Polynomial1D`

with ``degree=1``.

fitter : `astropy.modeling.fitting.Fitter`

Fitter for the given model to extract optimal coefficient values.

If `None`, will use `~astropy.modeling.fitting.LinearLSQFitter`.

power : int, optional

Power of values to fit.

min_datapoints : int, optional

Minimum number of data points required for fitting.

They can be built up with `time_func`.

Returns

-------

a : array-like

Fitted `~astropy.modeling.FittableModel` parameters.

Raises

------

AssertionError

Insufficient data points for fitting.

ModelsError

Invalid model or fitter.

"""

# Reset related attributes

self._power = power

self._cache_est = OrderedDict()

x_arr = np.array(list(six.iterkeys(self._cache_good)))

assert x_arr.size >= min_datapoints, \

'Requires {0} points but has {1}'.format(min_datapoints,

x_arr.size)

if model is None:

model = modeling.models.Polynomial1D(1)

elif not isinstance(model, modeling.core.Model):

raise modeling.fitting.ModelsError(

'{0} is not a model.'.format(model))

if fitter is None:

fitter = modeling.fitting.LinearLSQFitter()

elif not isinstance(fitter, modeling.fitting.Fitter):

raise modeling.fitting.ModelsError(

'{0} is not a fitter.'.format(fitter))

self._fit_func = fitter(

model, x_arr**power, list(six.itervalues(self._cache_good)))

return self._fit_func.parameters

def predict_time(self, arg):

"""Predict run time for given argument.

If prediction is already cached, cached value is returned.

Parameters

----------

arg : number

Input argument to predict run time for.

Returns

-------

t_est : float

Estimated run time for given argument.

Raises

------

AssertionError

No fitted data for prediction.

"""

if arg in self._cache_est:

t_est = self._cache_est[arg]

else:

assert self._fit_func is not None, 'No fitted data for prediction'

t_est = self._fit_func(arg**self._power)

self._cache_est[arg] = t_est

return t_est

def plot(self, xscale='linear', yscale='linear', xlabeltext='args',

save_as=''): # pragma: no cover

"""Plot prediction.

.. note:: Uses `matplotlib <http://matplotlib.sourceforge.net/>`_.

Parameters

----------

xscale, yscale : {'linear', 'log', 'symlog'}

Scaling for `matplotlib.axes.Axes`.

xlabeltext : str, optional

Text for X-label.

save_as : str, optional

Save plot as given filename.

Raises

------

AssertionError

Insufficient data for plotting.

"""

import matplotlib.pyplot as plt

# Actual data

x_arr = sorted(self._cache_good)

y_arr = np.array([self._cache_good[x] for x in x_arr])

assert len(x_arr) > 1, 'Insufficient data for plotting'

# Auto-ranging

qmean = y_arr.mean() * u.second

for cur_u in (u.minute, u.second, u.millisecond, u.microsecond,

u.nanosecond):

val = qmean.to(cur_u).value

if 1000 > val >= 1:

break

y_arr = (y_arr * u.second).to(cur_u).value

fig, ax = plt.subplots()

ax.plot(x_arr, y_arr, 'kx-', label='Actual')

# Fitted data

if self._fit_func is not None:

x_est = list(six.iterkeys(self._cache_est))

y_est = (np.array(list(six.itervalues(self._cache_est))) *

u.second).to(cur_u).value

ax.scatter(x_est, y_est, marker='o', c='r', label='Predicted')

x_fit = np.array(sorted(x_arr + x_est))

y_fit = (self._fit_func(x_fit**self._power) *

u.second).to(cur_u).value

ax.plot(x_fit, y_fit, 'b--', label='Fit')

ax.set_xscale(xscale)

ax.set_yscale(yscale)

ax.set_xlabel(xlabeltext)

ax.set_ylabel('Run time ({})'.format(cur_u.to_string()))

ax.set_title(self._funcname)

ax.legend(loc='best', numpoints=1)

plt.draw()

if save_as:

plt.savefig(save_as)