text:
import pandas as pd
import numpy as np
import statsmodels as sm
import statsmodels.api as smapi
import math
from pyqstrat.pq_utils import monotonically_increasing, infer_frequency
from pyqstrat.plot import TimeSeries, DateLine, Subplot, HorizontalLine, BucketedValues, Plot
import matplotlib as mpl
import matplotlib.figure as mpl_fig
from typing import Tuple, Sequence, Mapping, MutableMapping, Optional, Any, Callable, Dict

def compute_periods_per_year(timestamps: np.ndarray) -> float:
    """
    Computes trading periods per year for an array of numpy datetime64's.
        e.g. if most of the timestamps are separated by 1 day, will return 252.

    Args:
        timestamps: a numpy array of datetime64's

    >>> compute_periods_per_year(np.array(['2018-01-01', '2018-01-02', '2018-01-03', '2018-01-09'], dtype='M8[D]'))
    252.0
    >>> round(compute_periods_per_year(np.array(['2018-01-01 10:00', '2018-01-01 10:05', '2018-01-01 10:10'], dtype='M8[m]')), 2)
    72576.05
    """
    if not len(timestamps): return np.nan
    freq = infer_frequency(timestamps)
    return 252. / freq if freq != 0 else np.nan

def compute_amean(returns: np.ndarray, periods_per_year: int) -> float:
    '''
    Computes arithmetic mean of a return array, ignoring NaNs

    Args:
        returns: Represents returns at any frequency
        periods_per_year: Frequency of the returns, e.g. 252 for daily returns

    >>> compute_amean(np.array([0.003, 0.004, np.nan]), 252)
    0.882
    '''
    if not len(returns): return np.nan
    return np.nanmean(returns) * periods_per_year

def compute_num_periods(timestamps: np.ndarray, periods_per_year: float) -> float:
    '''
    Given an array of timestamps, we compute how many periods there are between the first and last element, where the length
        of a period is defined by periods_per_year.  For example, if there are 6 periods per year,
        then each period would be approx. 2 months long.

    Args:
        timestamps (np.ndarray of np.datetime64): a numpy array of returns, can contain nans
        periods_per_year: number of periods between first and last return

    >>> compute_num_periods(np.array(['2015-01-01', '2015-03-01', '2015-05-01'], dtype='M8[D]'), 6)
    2.0
    '''
    if not len(timestamps): return np.nan
    assert(monotonically_increasing(timestamps))
    fraction_of_year = (timestamps[-1] - timestamps[0]) / (np.timedelta64(1, 's') * 365 * 24 * 60 * 60)
    return round(fraction_of_year * periods_per_year)

def compute_gmean(timestamps: np.ndarray, returns: np.ndarray, periods_per_year: float) -> float:
    """
    Compute geometric mean of an array of returns

    Args:
        returns: a numpy array of returns, can contain nans
        periods_per_year: Used for annualizing returns

    >>> round(compute_gmean(np.array(['2015-01-01', '2015-03-01', '2015-05-01'], dtype='M8[D]'), np.array([0.001, 0.002, 0.003]), 252.), 6)
    0.018362
    """
    if not len(returns): return np.nan
    assert(len(returns) == len(timestamps))
    assert(isinstance(timestamps, np.ndarray) and isinstance(returns, np.ndarray))
    mask = np.isfinite(returns)
    timestamps = timestamps[mask]
    returns = returns[mask]
    num_periods = compute_num_periods(timestamps, periods_per_year)
    g_mean = ((1.0 + returns).prod())**(1.0 / num_periods)
    g_mean = np.power(g_mean, periods_per_year) - 1.0
    return g_mean

def compute_std(returns: np.ndarray) -> float:
    """ Computes standard deviation of an array of returns, ignoring nans """
    if not len(returns): return np.nan
    return np.nanstd(returns)

def compute_sortino(returns: np.ndarray, amean: float, periods_per_year: float) -> float:
    '''
    Note that this assumes target return is 0.

    Args:
        returns: a numpy array of returns
        amean: arithmetic mean of returns
        periods_per_year: number of trading periods per year

    >>> print(round(compute_sortino(np.array([0.001, -0.001, 0.002]), 0.001, 252), 6))
    0.133631
    '''
    if not len(returns) or not np.isfinite(amean) or periods_per_year <= 0: return np.nan
    returns = np.where((~np.isfinite(returns)), 0.0, returns)
    normalized_rets = np.where(returns > 0.0, 0.0, returns)
    sortino_denom = np.std(normalized_rets)
    sortino = np.nan if sortino_denom == 0 else amean / (sortino_denom * np.sqrt(periods_per_year))
    return sortino

def compute_sharpe(returns: np.ndarray, amean: float, periods_per_year: float) -> float:
    '''
    Note that this does not take into risk free returns so it's really a sharpe0, i.e. assumes risk free returns are 0

    Args:
        returns: a numpy array of returns
        amean: arithmetic mean of returns
        periods_per_year: number of trading periods per year

    >>> round(compute_sharpe(np.array([0.001, -0.001, 0.002]), 0.001, 252), 6)
    0.050508
    '''
    if not len(returns) or not np.isfinite(amean) or periods_per_year <= 0: return np.nan
    returns = np.where((~np.isfinite(returns)), 0.0, returns)
    s = np.std(returns)
    sharpe = np.nan if s == 0 else amean / (s * np.sqrt(periods_per_year))
    return sharpe

def compute_k_ratio(equity: np.ndarray, periods_per_year: int, halflife_years: float = None) -> float:
    '''
    Compute k-ratio (2013 or original versions by Lars Kestner). See https://papers.ssrn.com/sol3/papers.cfm?abstract_id=2230949
    We also implement a modification that allows higher weighting for more recent returns.

    Args:
        equity: a numpy array of the equity in your account
        periods_per_year: 252 for daily values
        halflife_years: If set, we use weighted linear regression to give less weight to older returns.
            In this case, we compute the original k-ratio which does not use periods per year or number of observations
            If not set, we compute the 2013 version of the k-ratio which weights k-ratio by sqrt(periods_per_year) / nobs

    Returns:
        weighted or unweighted k-ratio

    >>> np.random.seed(0)
    >>> t = np.arange(1000)
    >>> ret = np.random.normal(loc = 0.0025, scale = 0.01, size = len(t))
    >>> equity = (1 + ret).cumprod()
    >>> assert(math.isclose(compute_k_ratio(equity, 252, None), 3.888, abs_tol=0.001))
    >>> assert(math.isclose(compute_k_ratio(equity, 252, 0.5), 602.140, abs_tol=0.001))
    '''
    equity = equity[np.isfinite(equity)]
    equity = np.log(equity)
    t = np.arange(len(equity))
    if halflife_years:
        halflife = halflife_years * periods_per_year
        k = math.log(0.5) / halflife
        w = np.empty(len(equity), dtype=np.float)
        w = np.exp(k * t)
        w = w ** 2  # Statsmodels requires square of weights
        w = w[::-1]
        fit = sm.regression.linear_model.WLS(endog=equity, exog=t, weights=w, hasconst=False).fit()
        k_ratio = fit.params[0] / fit.bse[0]
    else:
        fit = smapi.OLS(endog=equity, exog=np.arange(len(equity)), hasconst=False).fit()
        k_ratio = fit.params[0] * math.sqrt(periods_per_year) / (fit.bse[0] * len(equity))

    return k_ratio

def compute_equity(timestamps: np.ndarray, starting_equity: float, returns: np.ndarray) -> np.ndarray:
    ''' Given starting equity, timestamps and returns, create a numpy array of equity at each date'''
    return starting_equity * np.cumprod(1. + returns)

def compute_rolling_dd(timestamps: np.ndarray, equity: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
    '''
    Compute numpy array of rolling drawdown percentage

    Args:
        timestamps: numpy array of datetime64
        equity: numpy array of equity
    '''
    assert(len(timestamps) == len(equity))
    if not len(timestamps): return np.array([], dtype='M8[ns]'), np.array([], dtype=np.float)
    s = pd.Series(equity, index=timestamps)
    rolling_max = s.expanding(min_periods=1).max()
    dd = np.where(s >= rolling_max, 0.0, -(s - rolling_max) / rolling_max)
    return timestamps, dd

def compute_maxdd_pct(rolling_dd: np.ndarray) -> float:
    '''Compute max drawdown percentage given a numpy array of rolling drawdowns, ignoring NaNs'''
    if not len(rolling_dd): return np.nan
    return np.nanmax(rolling_dd)

def compute_maxdd_date(rolling_dd_dates: np.ndarray, rolling_dd: np.ndarray) -> float:
    ''' Compute date of max drawdown given numpy array of timestamps, and corresponding rolling dd percentages'''
    if not len(rolling_dd_dates): return pd.NaT
    assert(len(rolling_dd_dates) == len(rolling_dd))
    return rolling_dd_dates[np.argmax(rolling_dd)]

def compute_maxdd_start(rolling_dd_dates: np.ndarray, rolling_dd: np.ndarray, mdd_date: np.datetime64) -> np.datetime64:
    '''Compute date when max drawdown starts, given numpy array of timestamps corresponding rolling dd
        percentages and date of the max draw down'''
    if not len(rolling_dd_dates) or pd.isnull(mdd_date): return pd.NaT
    assert(len(rolling_dd_dates) == len(rolling_dd))
    return rolling_dd_dates[(rolling_dd <= 0) & (rolling_dd_dates < mdd_date)][-1]

def compute_mar(returns: np.ndarray, periods_per_year: float, mdd_pct: float) -> float:
    '''Compute MAR ratio, which is annualized return divided by biggest drawdown since inception.'''
    if not len(returns) or np.isnan(mdd_pct) or mdd_pct == 0: return np.nan
    return np.mean(returns) * periods_per_year / mdd_pct

def compute_dates_3yr(timestamps: np.ndarray) -> np.ndarray:
    ''' Given an array of numpy datetimes, return those that are within 3 years of the last date in the array'''
    if not len(timestamps): return np.array([], dtype='M8[D]')
    last_date = timestamps[-1]
    d = pd.to_datetime(last_date)
    start_3yr = np.datetime64(d.replace(year=d.year - 3))
    return timestamps[timestamps > start_3yr]

def compute_returns_3yr(timestamps: np.ndarray, returns: np.ndarray) -> np.ndarray:
    '''Given an array of numpy datetimes and an array of returns, return those that are within 3 years
        of the last date in the datetime array '''
    if not len(timestamps): return np.array([], dtype=np.float)
    assert(len(timestamps) == len(returns))
    timestamps_3yr = compute_dates_3yr(timestamps)
    return returns[timestamps >= timestamps_3yr[0]]

def compute_rolling_dd_3yr(timestamps: np.ndarray, equity: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
    '''Compute rolling drawdowns over the last 3 years'''
    if not len(timestamps): return np.array([], dtype='M8[D]')
    last_date = timestamps[-1]
    d = pd.to_datetime(last_date)
    start_3yr = np.datetime64(d.replace(year=d.year - 3))
    equity = equity[timestamps >= start_3yr]
    timestamps = timestamps[timestamps >= start_3yr]
    return compute_rolling_dd(timestamps, equity)

def compute_maxdd_pct_3yr(rolling_dd_3yr: np.ndarray) -> float:
    '''Compute max drawdown percentage over the last 3 years'''
    return compute_maxdd_pct(rolling_dd_3yr)

def compute_maxdd_date_3yr(rolling_dd_3yr_timestamps: np.ndarray, rolling_dd_3yr: np.ndarray) -> np.datetime64:
    '''Compute max drawdown date over the last 3 years'''
    return compute_maxdd_date(rolling_dd_3yr_timestamps, rolling_dd_3yr)

def compute_maxdd_start_3yr(rolling_dd_3yr_timestamps: np.ndarray, rolling_dd_3yr: np.ndarray, mdd_date_3yr: np.datetime64) -> np.datetime64:
    '''Comput max drawdown start date over the last 3 years'''
    return compute_maxdd_start(rolling_dd_3yr_timestamps, rolling_dd_3yr, mdd_date_3yr)

def compute_calmar(returns_3yr: np.ndarray, periods_per_year: float, mdd_pct_3yr: float) -> float:
    '''Compute Calmar ratio, which is the annualized return divided by max drawdown over the last 3 years'''
    return compute_mar(returns_3yr, periods_per_year, mdd_pct_3yr)

def compute_bucketed_returns(timestamps: np.ndarray, returns: np.ndarray) -> Tuple[Sequence[int], Sequence[np.ndarray]]:
    '''
    Bucket returns by year

    Returns:
        A tuple with the first element being a list of years and the second a list of
            numpy arrays containing returns for each corresponding year
    '''
    assert(len(timestamps) == len(returns))
    if not len(timestamps): return np.array([], dtype=np.str), np.array([], dtype=np.float)
    s = pd.Series(returns, index=timestamps)
    years_list = []
    rets_list = []
    for year, rets in s.groupby(s.index.map(lambda x: x.year)):
        years_list.append(year)
        rets_list.append(rets.values)

    return years_list, rets_list

def compute_annual_returns(timestamps: np.ndarray, returns: np.ndarray, periods_per_year: float) -> Tuple[np.ndarray, np.ndarray]:
    '''Takes the output of compute_bucketed_returns and returns geometric mean of returns by year

    Returns:
        A tuple with the first element being an array of years (integer) and the second element
        an array of annualized returns for those years

    '''
    assert(len(timestamps) == len(returns) and periods_per_year > 0)
    if not len(timestamps): return np.array([], dtype=np.str), np.array([], dtype=np.float)
    df = pd.DataFrame({'ret': returns, 'timestamp': timestamps})
    years = []
    gmeans = []
    for k, g in df.groupby(df.timestamp.map(lambda x: x.year)):
        years.append(k)
        gmeans.append(compute_gmean(g.timestamp.values, g.ret.values, periods_per_year))
    return np.array(years), np.array(gmeans)

class Evaluator:
    """You add functions to the evaluator that are dependent on the outputs of other functions.
    The evaluator will call these functions in the right order
    so dependencies are computed first before the functions that need their output.
    You can retrieve the output of a metric using the metric member function

    >>> evaluator = Evaluator(initial_metrics={'x': np.array([1, 2, 3]), 'y': np.array([3, 4, 5])})
    >>> evaluator.add_metric('z', lambda x, y: sum(x, y), dependencies=['x', 'y'])
    >>> evaluator.compute()
    >>> evaluator.metric('z')
    array([ 9, 10, 11])
    """
    def __init__(self, initial_metrics: Dict[str, Any]) -> None:
        """Inits Evaluator with a dictionary of initial metrics that are used to compute subsequent metrics

        Args:
            initial_metrics: a dictionary of string name -> metric.  metric can be any object including a scalar,
                an array or a tuple
        """
        assert(type(initial_metrics) == dict)
        self.metric_values: Dict[str, Any] = initial_metrics.copy()
        self._metrics: MutableMapping[str, Tuple[Callable, Sequence[str]]] = {}

    def add_metric(self, name: str, func: Callable, dependencies: Sequence[str]) -> None:
        self._metrics[name] = (func, dependencies)

    def compute(self, metric_names: Sequence[str] = None) -> None:
        '''Compute metrics using the internal dependency graph

        Args:
            metric_names: an array of metric names.  If not passed in, evaluator will compute and store all metrics
        '''

        if metric_names is None: metric_names = list(self._metrics.keys())
        for metric_name in metric_names:
            self.compute_metric(metric_name)

    def compute_metric(self, metric_name: str) -> None:
        '''
        Compute and store a single metric:

        Args:
            metric_name: string representing the metric to compute
        '''
        func, dependencies = self._metrics[metric_name]
        for dependency in dependencies:
            if dependency not in self.metric_values:
                self.compute_metric(dependency)
        dependency_values = {k: self.metric_values[k] for k in dependencies}

        values = func(**dependency_values)

        self.metric_values[metric_name] = values

    def metric(self, metric_name: str) -> Any:
        '''Return the value of a single metric given its name'''
        return self.metric_values[metric_name]

    def metrics(self) -> Mapping[str, Any]:
        '''Return a dictionary of metric name -> metric value'''
        return self.metric_values

def handle_non_finite_returns(timestamps: np.ndarray,
                              rets: np.ndarray,
                              leading_non_finite_to_zeros: bool,
                              subsequent_non_finite_to_zeros: bool) -> Tuple[np.ndarray, np.ndarray]:
    '''
    >>> np.set_printoptions(formatter={'float': '{: .6g}'.format})
    >>> timestamps = np.arange(np.datetime64('2019-01-01'), np.datetime64('2019-01-07'))
    >>> rets = np.array([np.nan, np.nan, 3, 4, np.nan, 5])
    >>> handle_non_finite_returns(timestamps, rets, leading_non_finite_to_zeros = False, subsequent_non_finite_to_zeros = True)
    (array(['2019-01-03', '2019-01-04', '2019-01-05', '2019-01-06'], dtype='datetime64[D]'), array([ 3,  4,  0,  5]))
    >>> handle_non_finite_returns(timestamps, rets, leading_non_finite_to_zeros = True, subsequent_non_finite_to_zeros = False)
    (array(['2019-01-01', '2019-01-02', '2019-01-03', '2019-01-04', '2019-01-06'], dtype='datetime64[D]'), array([ 0,  0,  3,  4,  5]))
    >>> handle_non_finite_returns(timestamps, rets, leading_non_finite_to_zeros = False, subsequent_non_finite_to_zeros = False)
    (array(['2019-01-01', '2019-01-02', '2019-01-03', '2019-01-04', '2019-01-06'], dtype='datetime64[D]'), array([ 0,  0,  3,  4,  5]))
    >>> rets = np.array([1, 2, 3, 4, 4.5,  5])
    >>> handle_non_finite_returns(timestamps, rets, leading_non_finite_to_zeros = False, subsequent_non_finite_to_zeros = True)
    (array(['2019-01-01', '2019-01-02', '2019-01-03', '2019-01-04', '2019-01-05', '2019-01-06'],
        dtype='datetime64[D]'), array([ 1,  2,  3,  4,  4.5,  5]))
    '''

    first_non_nan_index = np.ravel(np.nonzero(~np.isnan(rets)))
    if len(first_non_nan_index):
        first_non_nan_index = first_non_nan_index[0]
    else:
        first_non_nan_index = -1

    if first_non_nan_index > 0 and first_non_nan_index < len(rets):
        if leading_non_finite_to_zeros:
            rets[:first_non_nan_index] = np.nan_to_num(rets[:first_non_nan_index])
        else:
            timestamps = timestamps[first_non_nan_index:]
            rets = rets[first_non_nan_index:]

    if subsequent_non_finite_to_zeros:
        rets = np.nan_to_num(rets)
    else:
        timestamps = timestamps[np.isfinite(rets)]
        rets = rets[np.isfinite(rets)]

    return timestamps, rets

def compute_return_metrics(timestamps: np.ndarray,
                           rets: np.ndarray,
                           starting_equity: float,
                           leading_non_finite_to_zeros: bool = False,
                           subsequent_non_finite_to_zeros: bool = True) -> Evaluator:
    '''
    Compute a set of common metrics using returns (for example, of an instrument or a portfolio)

    Args:
        timestamps (np.array of datetime64): Timestamps for the returns
        rets (nd.array of float): The returns, use 0.01 for 1%
        starting_equity (float): Starting equity value in your portfolio
        leading_non_finite_to_zeros (bool, optional): If set, we replace leading nan, inf, -inf returns with zeros.
            For example, you may need a warmup period for moving averages.  Default False
        subsequent_non_finite_to_zeros (bool, optional): If set, we replace any nans that follow the first non nan value with zeros.
            There may be periods where you have no prices but removing these returns would result in incorrect annualization.
            Default True

    Returns:
        An Evaluator object containing computed metrics off the returns passed in.
        If needed, you can add your own metrics to this object based on the values of existing metrics and recompute the Evaluator.
        Otherwise, you can just use the output of the evaluator using the metrics function.

    >>> timestamps = np.array(['2015-01-01', '2015-03-01', '2015-05-01', '2015-09-01'], dtype='M8[D]')
    >>> rets = np.array([0.01, 0.02, np.nan, -0.015])
    >>> starting_equity = 1.e6
    >>> ev = compute_return_metrics(timestamps, rets, starting_equity)
    >>> metrics = ev.metrics()
    >>> assert(round(metrics['gmean'], 6) == 0.021061)
    >>> assert(round(metrics['sharpe'], 6) == 0.599382)
    >>> assert(all(metrics['returns_3yr'] == np.array([0.01, 0.02, 0, -0.015])))
    '''

    assert(starting_equity > 0.)
    assert(type(rets) == np.ndarray and rets.dtype == np.float64)
    assert(type(timestamps) == np.ndarray and np.issubdtype(timestamps.dtype, np.datetime64) and monotonically_increasing(timestamps))

    timestamps, rets = handle_non_finite_returns(timestamps, rets, leading_non_finite_to_zeros, subsequent_non_finite_to_zeros)

    ev = Evaluator({'timestamps': timestamps, 'returns': rets, 'starting_equity': starting_equity})
    ev.add_metric('periods_per_year', compute_periods_per_year, dependencies=['timestamps'])
    ev.add_metric('amean', compute_amean, dependencies=['returns', 'periods_per_year'])
    ev.add_metric('std', compute_std, dependencies=['returns'])
    ev.add_metric('up_periods', lambda returns: len(returns[returns > 0]), dependencies=['returns'])
    ev.add_metric('down_periods', lambda returns: len(returns[returns < 0]), dependencies=['returns'])
    ev.add_metric('up_pct',
                  lambda up_periods, down_periods: up_periods * 1.0 / (up_periods + down_periods) if (up_periods + down_periods) != 0 else np.nan,
                  dependencies=['up_periods', 'down_periods'])
    ev.add_metric('gmean', compute_gmean, dependencies=['timestamps', 'returns', 'periods_per_year'])
    ev.add_metric('sharpe', compute_sharpe, dependencies=['returns', 'periods_per_year', 'amean'])
    ev.add_metric('sortino', compute_sortino, dependencies=['returns', 'periods_per_year', 'amean'])
    ev.add_metric('equity', compute_equity, dependencies=['timestamps', 'starting_equity', 'returns'])
    ev.add_metric('k_ratio', compute_k_ratio, dependencies=['equity', 'periods_per_year'])
    ev.add_metric('k_ratio_weighted', lambda equity, periods_per_year: compute_k_ratio(equity, periods_per_year, 3),
                  dependencies=['equity', 'periods_per_year'])

    # Drawdowns
    ev.add_metric('rolling_dd', compute_rolling_dd, dependencies=['timestamps', 'equity'])
    ev.add_metric('mdd_pct', lambda rolling_dd: compute_maxdd_pct(rolling_dd[1]), dependencies=['rolling_dd'])
    ev.add_metric('mdd_date', lambda rolling_dd: compute_maxdd_date(rolling_dd[0], rolling_dd[1]), dependencies=['rolling_dd'])
    ev.add_metric('mdd_start', lambda rolling_dd, mdd_date: compute_maxdd_start(rolling_dd[0], rolling_dd[1], mdd_date),
                  dependencies=['rolling_dd', 'mdd_date'])
    ev.add_metric('mar', compute_mar, dependencies=['returns', 'periods_per_year', 'mdd_pct'])

    ev.add_metric('timestamps_3yr', compute_dates_3yr, dependencies=['timestamps'])
    ev.add_metric('returns_3yr', compute_returns_3yr, dependencies=['timestamps', 'returns'])

    ev.add_metric('rolling_dd_3yr', compute_rolling_dd_3yr, dependencies=['timestamps', 'equity'])
    ev.add_metric('mdd_pct_3yr', lambda rolling_dd_3yr: compute_maxdd_pct_3yr(rolling_dd_3yr[1]), dependencies=['rolling_dd_3yr'])
    ev.add_metric('mdd_date_3yr', lambda rolling_dd_3yr: compute_maxdd_date_3yr(rolling_dd_3yr[0], rolling_dd_3yr[1]),
                  dependencies=['rolling_dd_3yr'])
    ev.add_metric('mdd_start_3yr', lambda rolling_dd_3yr, mdd_date_3yr:
                  compute_maxdd_start_3yr(rolling_dd_3yr[0], rolling_dd_3yr[1], mdd_date_3yr),
                  dependencies=['rolling_dd_3yr', 'mdd_date_3yr'])
    ev.add_metric('calmar', compute_calmar, dependencies=['returns_3yr', 'periods_per_year', 'mdd_pct_3yr'])

    ev.add_metric('annual_returns', compute_annual_returns, dependencies=['timestamps', 'returns', 'periods_per_year'])
    ev.add_metric('bucketed_returns', compute_bucketed_returns, dependencies=['timestamps', 'returns'])

    ev.compute()
    return ev

def display_return_metrics(metrics: Mapping[str, Any], float_precision: int = 3) -> pd.DataFrame:
    '''
    Creates a dataframe making it convenient to view the output of the metrics obtained using the compute_return_metrics function.

    Args:
        float_precision: Change if you want to display floats with more or less significant figures than the default,
            3 significant figures.
    Returns:
        A one row dataframe with formatted metrics.
    '''
    from IPython.core.display import display

    _metrics = {}
    cols = ['gmean', 'amean', 'std', 'shrp', 'srt', 'k', 'calmar', 'mar', 'mdd_pct', 'mdd_start', 'mdd_date', 'dd_3y_pct',
            'up_periods', 'down_periods', 'up_pct', 'mdd_start_3yr', 'mdd_date_3yr']

    translate = {'shrp': 'sharpe', 'srt': 'sortino', 'dd_3y_pct': 'mdd_pct_3yr', 'k': 'k_ratio'}
    for col in cols:
        key = col
        if col in translate: key = translate[col]
        _metrics[col] = metrics[key]

    _metrics['mdd_dates'] = f'{str(metrics["mdd_start"])[:10]}/{str(metrics["mdd_date"])[:10]}'
    _metrics['up_dwn'] = f'{metrics["up_periods"]}/{metrics["down_periods"]}/{metrics["up_pct"]:.3g}'
    _metrics['dd_3y_timestamps'] = f'{str(metrics["mdd_start_3yr"])[:10]}/{str(metrics["mdd_date_3yr"])[:10]}'

    years = metrics['annual_returns'][0]
    ann_rets = metrics['annual_returns'][1]
    for i, year in enumerate(years):
        _metrics[str(year)] = ann_rets[i]

    format_str = '{:.' + str(float_precision) + 'g}'

    for k, v in _metrics.items():
        if isinstance(v, np.float) or isinstance(v, float):
            _metrics[k] = format_str.format(v)

    cols = ['gmean', 'amean', 'std', 'shrp', 'srt', 'k', 'calmar', 'mar', 'mdd_pct', 'mdd_dates', 'dd_3y_pct', 'dd_3y_timestamps', 'up_dwn'] + [
        str(year) for year in sorted(years)]

    df = pd.DataFrame(index=[''])
    for metric_name, metric_value in _metrics.items():
        df.insert(0, metric_name, metric_value)
    df = df[cols]

    display(df)
    return df

def plot_return_metrics(metrics: Mapping[str, Any], title: str = None) -> Optional[Tuple[mpl_fig.Figure, mpl.axes.Axes]]:
    '''
    Plot equity, rolling drawdowns and and a boxplot of annual returns given the output of compute_return_metrics.
    '''
    timestamps = metrics['timestamps']
    equity = metrics['equity']
    equity = TimeSeries('equity', timestamps=timestamps, values=equity)
    mdd_date, mdd_start = metrics['mdd_start'], metrics['mdd_date']
    mdd_date_3yr, mdd_start_3yr = metrics['mdd_start_3yr'], metrics['mdd_date_3yr']
    drawdown_lines = [DateLine(name='max dd', date=mdd_start, color='red'),
                      DateLine(date=mdd_date, color='red'),
                      DateLine(name='3y dd', date=mdd_start_3yr, color='orange'),
                      DateLine(date=mdd_date_3yr, color='orange')]
    equity_subplot = Subplot(equity, ylabel='Equity', height_ratio=0.6, log_y=True, y_tick_format='${x:,.0f}',
                             date_lines=drawdown_lines, horizontal_lines=[HorizontalLine(metrics['starting_equity'], color='black')])

    rolling_dd = TimeSeries('drawdowns', timestamps=metrics['rolling_dd'][0], values=metrics['rolling_dd'][1])
    zero_line = HorizontalLine(y=0, color='black')
    dd_subplot = Subplot(rolling_dd, ylabel='Drawdowns', height_ratio=0.2, date_lines=drawdown_lines, horizontal_lines=[zero_line])

    years = metrics['bucketed_returns'][0]
    ann_rets = metrics['bucketed_returns'][1]
    ann_ret = BucketedValues('annual returns', bucket_names=years, bucket_values=ann_rets)
    ann_ret_subplot = Subplot(ann_ret, ylabel='Annual Returns', height_ratio=0.2, horizontal_lines=[zero_line])

    plt = Plot([equity_subplot, dd_subplot, ann_ret_subplot], title=title)
    return plt.draw()

def test_evaluator() -> None:
    from datetime import datetime, timedelta
    np.random.seed(10)
    timestamps = np.arange(datetime(2018, 1, 1), datetime(2018, 3, 1), timedelta(days=1))
    rets = np.random.normal(size=len(timestamps)) / 1000
    starting_equity = 1.e6

    ev = compute_return_metrics(timestamps, rets, starting_equity)
    display_return_metrics(ev.metrics())
    plot_return_metrics(ev.metrics())

    assert(round(ev.metric('sharpe'), 6) == 2.932954)
    assert(round(ev.metric('sortino'), 6) == 5.690878)
    assert(ev.metric('annual_returns')[0] == [2018])
    assert(round(ev.metric('annual_returns')[1][0], 6) == [0.063530])
    assert(ev.metric('mdd_start') == np.datetime64('2018-01-19'))
    assert(ev.metric('mdd_date') == np.datetime64('2018-01-22'))

if __name__ == "__main__":
    test_evaluator()
    import doctest
    doctest.testmod(optionflags=doctest.NORMALIZE_WHITESPACE)

max_stars_repo_name: alexanu/pyqstrat

max_issues_repo_name: alexanu/pyqstrat

max_forks_repo_name: alexanu/pyqstrat

reason: import numpy,import statsmodels


text:
#!/usr/bin/env python

from pandas  import *
from numpy  import *
from djeval import *
import csv, code
import cPickle as pickle
from sklearn.externals import joblib

GAMELIMIT=int(sys.argv[2])

# Takes ~60 seconds on 100k games

def shell():
    vars = globals()
    vars.update(locals())
    shell = code.InteractiveConsole(vars)
    shell.interact()

def writecol(foo, comma=True):
    if foo is not None:
        if type(foo) == np.float64:
            outfd.write('%s' % round(foo,3))
        else:
            outfd.write('%s' % foo)
    if comma:
        outfd.write(',')

msg("Hi! Reading eheaders")
eheaders_filename = '/data/eheaders.p'
eheaders_file = open(eheaders_filename, 'r')
eheaders = pickle.load(eheaders_file)
elos = eheaders['elos']
result = eheaders['result']
checkmate = eheaders['checkmate']
openings = eheaders['openings']
ocount = eheaders['opening_count']
timecontrols = eheaders['timecontrols']

msg("Writing crunched movescores to %s" % sys.argv[1])
outfd = open(sys.argv[1], 'w')

msg("Hi! Reading movescores into memory, using kaggle-supplied scores as a backstop")
rows = {}
for scorefile_name in ['/data/movescores.csv', '/data/stockfish.csv']:
    stockfish_scores = open(scorefile_name)
    stockfish_reader = csv.reader(stockfish_scores, delimiter=',')
    for row in stockfish_reader:
        if row[0] == 'Event':
            continue
        gamenum = int(row[0])
        if gamenum not in rows:
            rows[gamenum] = row

def was_matecreated(prev, next):
  return (prev > -500 and next < -1000)

def was_matedestroyed(prev, next):
  return (prev > 1000 and next < 500)

writecol('gamenum')
writecol('side')
writecol('meanerror')
writecol('stdeverror')
writecol('q_error_one')
writecol('q_error_two')
writecol('meanecho')
writecol('blunderrate')
writecol('perfectrate')
writecol('gameoutcome')
writecol('gamelength')
writecol('grit')
writecol('won_by_checkmate')
writecol('lost_by_checkmate')
writecol('my_final_equity')
writecol('stdevpos')
writecol('pos_fft_1')
writecol('mate_created')
writecol('mate_destroyed')
writecol('early_lead', comma=False)
outfd.write('\n')

msg("Hi! Running through movescores")
for row in rows.values():

  gamenum = int(row[0])
  movescores = {}
  matecreated = {}
  matedestroyed = {}
  movescores[1] = []
  movescores[-1] = []
  moverecho = {}
  moverecho[1] = []
  moverecho[-1] = []
  grit = {}
  grit[1] = 0
  grit[-1] = 0
  lead_established = False
  position_scores = []

  if gamenum > GAMELIMIT:
    break

  if gamenum % 500 == 0:
    msg("hi doing gamenum %i" % gamenum)

  strscores = row[1].split(' ')
  side = 1
  try:
      last_equity = int(strscores[0])
  except ValueError:
      sys.stderr.write('Couldnt parse row, skipping: %s\n' % row)
      next
  last_gain = 0
  movenum = 0
  num_moves_while_losing = 0
  matecreated[1] = False
  matecreated[-1] = False
  matedestroyed[1] = False
  matedestroyed[-1] = False
  early_lead = None

  for strscore in strscores[1:]:
    # only for stockfish.csv
    if (strscore == 'NA') or (strscore == ''):
      score = 0
    else:
      score = int(strscore)

    position_scores.append(score)
    whitegain = None
    movergain = None

    if not lead_established and abs(score) > 100:
      lead_established = True
      early_lead = movenum

    if last_equity is not None:
      whitegain = score - last_equity
      movergain = whitegain * side
      if last_equity * side < -300:
        grit[side] = grit[side] + 1
      if last_gain == 0:
        echo = 0
      else:
        echo = max(0, min(1, movergain / last_gain))
      if abs(last_equity) < 30000:
        movescores[side].append(movergain)
        moverecho[side].append(echo)
      if was_matecreated(side * last_equity, side * score):
        matecreated[side] = True
      if was_matedestroyed(side * last_equity, side * score):
        matedestroyed[side] = True

    mover_score = last_equity * side

    last_equity = score
    last_gain = movergain
    side = side * -1
    movenum = movenum + 1

  for side in [-1, 1]:
    clippederror = clip(movescores[side], -150, 0)
    if len(clippederror) == 0:
        clippederror = array([-15])
        moverecho[side] = array([0.1])

    writecol(gamenum)
    writecol(side)
    writecol(mean(clippederror))
    writecol(clippederror.std())
    writecol(percentile(clippederror, 25))
    writecol(percentile(clippederror, 10))
    writecol(mean(moverecho[side]))
    writecol(mean(clippederror < -100))
    writecol(mean(clippederror == 0))
    writecol(result[gamenum] * side)
    writecol(len(strscores))
    writecol(grit[side])
    writecol((1 == result[gamenum] * side) & checkmate[gamenum])
    writecol((-1 == result[gamenum] * side) & checkmate[gamenum])
    if strscores[-1] != 'NA' and strscores[-1] != '':
      final_equity = side * int(strscores[-1])
    else:
      final_equity = 0
    writecol(final_equity)
    writecol(clip(position_scores, -500, 500).std())
    writecol(fft.fft(position_scores, 5).real[1])
    writecol(matecreated[side])
    writecol(matedestroyed[side])
    writecol(early_lead, comma=False)
    outfd.write('\n')

ext: py

max_stars_repo_name: Keesiu/meta-kaggle

max_issues_repo_name: Keesiu/meta-kaggle

max_forks_repo_name: Keesiu/meta-kaggle

reason: from numpy


text:
from astropy import units as u
from astropy.cosmology import Planck15
import numpy as np
from scipy.integrate import cumtrapz, quad, simps, trapz
from time import time

try:
    import pyccl as ccl
    has_ccl = True
except ImportError:
    has_ccl = False

from .helpers.cosmology import BaseCosmo
from .helpers.decorators import array, inMpc
from .helpers.lensing import BaseLensing

class Profile(BaseLensing):
    """Profile object

    All profiles should inherit from ``Profile``

    Defining your own profile is very simple. As an example, let's
    define a simple power-law profile with two free parameters, the
    normalization and the slope:

    .. math::

        f(r, a, b) = a * r**b

    ::

        class PowerLaw(Profile):

            def __init__(self, norm, slope, **kwargs):
                self._set_shape(norm*slope)
                super().__init__(**kwargs)
                self.norm = norm
                self.slope = slope

            @array
            def profile(self, r):
                return self.norm * r**self.slope

    That's it! The ``__init__()`` method needs only two lines of code
    (in addition to attribute definitions). The last line is necessary
    to allow ``profiley`` to automatically handle arbitrary shapes,
    through the definition of a ``_shape`` attribute. Note that
    ``_set_shape`` takes only one argument (besides ``self``) - the
    *product* of the class arguments. That is, if the arguments are
    arrays, their dimensions must be such that a product can be carried
    out without any manipulation.

    Profile projections
    -------------------

    If the projection of this profile is analytical, any or all of the
    following methods can also be specified: ::

        projected(self, R)
        projected_cumulative(self, R)
        projected_excess(self, R)
        offset_profile(self, R, Roff)
        offset_projected(self, R, Roff)
        offset_projected_cumulative(self, R, Roff)
        offset_projected_excess(self, R, Roff)

    If it does not have analytical expressions, these methods will also
    exist, but they will be calculated numerically, so they may be
    somewhat slower depending on the precision required.

    Cosmology
    ---------

    All ``Profile`` objects contain all cosmological information with
    they have been initialized through the ``self.cosmo`` attribute,
    which can be any ``astropy.cosmology.FLRW`` object.

    """
    def __init__(self, z=0, overdensity=500,
                 los_loglimit=6, nsamples_los=200, resampling=20,
                 logleft=-10, left_samples=100, **kwargs):
        """Initialize a profile object

        Optional arguments
        ------------------
        z : float or ndarray of floats
            redshift
        overdensity : int or float
            overdensity with respect to the background (does not apply
            to all Profile children; see each specific class for
            details)
        """
        super().__init__(z, **kwargs)
        # check overdensity
        self._assert_overdensity(overdensity)
        self.overdensity = overdensity
        # for numerical integration -- perhaps these could be passed
        # in a single dictionary
        self.los_loglimit = los_loglimit
        self.nsamples_los = nsamples_los
        self.resampling = resampling
        self.logleft = logleft
        self.left_samples = left_samples
        # empty init
        self.__dimensions = None

    @property
    def _one(self):
        if self.__one is None:
            self.__one = u.dimensionless_unscaled
        return self.__one

    @property
    def _dimensions(self):
        if self.__dimensions is None:
            self.__dimensions = tuple([1] * len(self.shape))
        return self.__dimensions

    @property
    def shape(self):
        if not hasattr(self, '_shape'):
            msg = 'attribute shape does not exist. Please make sure' \
                  ' to call ``self._set_shape`` in ``__init__``'
            raise AttributeError(msg)
        return self._shape

    ### private methods ###

    def _assert_overdensity(self, overdensity):
        assert not np.iterable(overdensity), \
            'parameter overdensity must be a scalar'
        try:
            overdensity / 1
        except TypeError as err:
            raise TypeError('parameter overdensity must be a number') from err
        if overdensity <= 0:
            raise ValueError(
                f'overdensity must be positive; received {overdensity}')
        return

    def _define_array(self, x):
        if not np.iterable(x):
            return x * np.ones(self._shape)
        return x

    def _set_shape(self, args_product):
        if hasattr(self, 'shape'):
            msg = 'attribute shape already set, cannot be overwritten'
            raise ValueError(msg)
        if np.iterable(args_product):
            self._shape = args_product.shape
        else:
            self._shape = (1,)

    ### methods ###

    @inMpc
    @array
    def projected(self, R: np.ndarray, log_rmin=-10, log_rmax=6,
                  integral_samples=200):
        """Line of sight projected profile, calculated numerically

        Parameters
        ----------
        R : np.ndarray
            positions at which to calculate the projected profile

        Optional arguments
        ------------------
        log_rmin, log_rmax : float
            lower and upper limits for logspace resampling for integration
        integral_samples : int
            number of samples to generate for Simpson-rule integration
            of the projected profile

        Notes on numerical integration
        ------------------------------
         -The default values for the integration parameters give
            numerical errors well below 0.1% over the range
            R=[1e-5,100] Mpc, when comparing the numerical and
            analytical implementations for an NFW profile (the
            former can be obtained by defining a GNFW profile with
            default kwargs)
        """
        assert log_rmin < log_rmax, \
            'argument log_rmin must be larger than log_rmax, received' \
            f' {log_rmin,log_rmax}'
        assert integral_samples // 1 == integral_samples, \
            'argument integral_samples must be int, received' \
            f' {integral_samples} ({type(integral_samples)})'
        R_los = np.logspace(log_rmin, log_rmax, integral_samples)[:,None]
        R = np.transpose(
            [np.hypot(*np.meshgrid(R_los[:,i], R[:,0]))
             for i in range(R_los.shape[1])],
            axes=(1,2,0))
        return 2 * simps(self.profile(R), R_los[None], axis=1)

    @inMpc
    @array
    def projected_cumulative(self, R: np.ndarray, log_rmin: float=-10,
                             left_samples: int=100, resampling: int=20,
                             **kwargs):
        """Cumulative projected profile within R, calculated
        numerically

        Parameters
        ----------
        R : np.ndarray
            positions at which to calculate the projected profile

        Optional arguments
        ------------------
        log_rmin : float
            lower limit for logspace resampling for integration. The
            same value will be passed to ``self.projected``
        resampling : int
            number of samples into which each R-interval in the
            data will be re-sampled. For instance, if two adjacent
            data points are at Rbin=0.1,0.2 then for the integration
            they will be replaced by
                newRbin = np.logspace(np.log10(0.1), np.log10(0.2),
                                      resampling, endpoint=False)
            (the endpoint will be added when sampling the following bin)
        left_samples : int
            number of samples to use between log_rmin and the first
            value of R, with a logarithmic sampling

        Additional arguments will be passed to ``self.projected``

        Notes on numerical integration
        ------------------------------
         -The default values for the integration parameters give
            numerical errors well below 0.1% over the range
            R=[1e-5,100] Mpc, when comparing the numerical and
            analytical implementations for an NFW profile (the
            former can be obtained by defining a GNFW profile with
            default kwargs)
        """
        assert isinstance(left_samples, (int,np.integer)), \
            'argument left_samples must be int, received' \
            f' {left_samples} ({type(left_samples)})'
        assert isinstance(resampling, (int,np.integer)), \
            'argument resampling must be int, received' \
            f' {resampling} ({type(resampling)})'
        # for convenience
        logR = np.log10(R[:,0])
        # resample R
        Ro = np.vstack(
            [np.zeros(R.shape[1]),
             np.logspace(log_rmin, logR[0], left_samples, endpoint=False)[:,None],
             np.concatenate(
                [np.logspace(logR[i-1], logR[i], resampling, endpoint=False)
                 for i in range(1, R.shape[0])])[:,None],
             R.max()*np.ones(R.shape[1])]
            )
        j = np.arange(1+left_samples, Ro.shape[0], resampling)
        integ = cumtrapz(
            Ro*self.projected(Ro, log_rmin=log_rmin, **kwargs),
            Ro, initial=0, axis=0)
        return 2 * integ[j] / R**2

    def projected_excess(self, R: np.ndarray, log_rmin=-10, log_rmax=6,
                         integral_samples=200,
                         left_samples=100, resampling=20):
        """Cumulative projected profile file excess at projected
        distance(s) R, defined as

            projected_excess(R) = projected_cumulative(R) - projected(R)

        This profile is most commonly used as the galaxy weak lensing
        *shear* observable, :math:`\gamma` where the projected excess
        is referred to as the *excess surface density* (ESD or
        :math:`\Delta\Sigma`),

        .. math::

            \Delta\Sigma(R) = \gamma\Sigma_\mathrm{c}

        where :math:`\Sigma_\mathrm{c}` is the critical surface density

        Parameters
        ----------
        R : float or array of float
            projected distance(s)

        Optional arguments are passed to either ``self.projected`` or
        ``self.projected_cumulative``
        """
        s1 = self.projected_cumulative(
            R, log_rmin=log_rmin, left_samples=left_samples,
            resampling=resampling, log_rmax=log_rmax,
            integral_samples=integral_samples)
        s2 = self.projected(R, log_rmin=log_rmin, log_rmax=log_rmax,
            integral_samples=integral_samples)
        return s1 - s2

    def offset(self, func, R, Roff, theta_samples=360, weights=None,
               **kwargs):
        """Calcuate any profile with a reference point different
        from its center

        Parameters
        ----------
        func : callable
            the funcion to calculate
        R : np.ndarray, shape (N,)
            radii at which to calculate the offset surface density
        Roff : np.ndarray, shape (M,)
            offsets with respect to the profile center

        Optional parameters
        -------------------
        theta_samples : int
            number of samples for the angular integral from 0 to 2*pi
        weights : array of floats, shape (M,)
            weights to apply to each profile corresponding to every
            value of ``Roff``. See ``Returns`` below
        kwargs : dict
            arguments to pass to ``func``

        Returns
        -------
        offset : np.ndarray,
            offset profile. The shape of the array depends on whether
            the ``weights`` argument is specified: if *not* specified
            (default), then
            .. code-block::

                shape: (M,N,*self.shape)

            if ``weights`` is provided, then the first axis will be
            weight-averaged over so that
            .. code-block::

                shape: (N,*self.shape)

        """
        if not isinstance(theta_samples, (int,np.integer)):
            raise TypeError(
                'argument theta_samples must be int, received' \
                f' {theta_samples} ({type(theta_samples)})')
        if not np.iterable(Roff):
            Roff = np.array([Roff])
        assert len(Roff.shape) == 1, 'argument Roff must be 1d'
        if weights is not None:
            if weights.size != Roff.size:
                msg = 'weights must have the same size as Roff,' \
                      f' received {weights.size}, {Roff.size},' \
                      ' respectively.'
                raise ValueError(msg)

        # can't get this to work using the @array decorator
        R = R.reshape((R.size,*self._dimensions))
        Roff = Roff.reshape((Roff.size,*self._dimensions,1,1))
        theta = np.linspace(0, 2*np.pi, theta_samples)
        theta1 = theta.reshape((theta_samples,*self._dimensions,1))
        x = (Roff**2 + R**2 + 2*R*Roff*np.cos(theta1))**0.5
        # looping slower but avoids memory issues
        # generator for the function calls beforehand makes it a little faster
        f = (func(i, **kwargs) for i in x)
        off = np.array([trapz(fi, theta, axis=0) for fi in f])

        if weights is not None:
            # create a slice so we can multiply by weights
            # along the first axis
            s_ = [None] * off.ndim
            s_[0] = slice(None)
            Roff = np.squeeze(Roff)
            off = trapz(weights[tuple(s_)]*off, Roff, axis=0) \
                / trapz(weights, Roff)
        return off / (2*np.pi)

    def offset_profile(self, R, Roff, **kwargs):
        """Alias for ``offset(profile, R, Roff, **kwargs)``"""
        return self.offset(self.profile, R, Roff)

    def offset_projected(self, R, Roff, **kwargs):
        """Alias for ``offset(projected, R, Roff, **kwargs)``"""
        return self.offset(self.projected, R, Roff, **kwargs)

    def offset_projected_cumulative(self, R, Roff, **kwargs):
        """Alias for ``offset(projected_cumulative, R, Roff,
        **kwargs)``"""
        return self.offset(self.projected_cumulative, R, Roff, **kwargs)

    def offset_projected_excess(self, R, Roff, **kwargs):
        """Alias for ``offset(projected_excess, R, Roff,
        **kwargs)``"""
        return self.offset(self.projected_excess, R, Roff, **kwargs)

    def _fourier(self, rmax=10, dr=0.1):
        """This is not working yet! Might just need to fall back to quad"""
        r = np.arange(dr, rmax, dr)
        f = self.profile(r)
        # compute Fourier transform by numpy's FFT function
        g = np.fft.fft(f)
        print('g =', g.shape)
        # frequency normalization factor is 2*np.pi/dt
        k = np.fft.fftfreq(f.size)*2*np.pi/dr
        # in order to get a discretisation of the continuous
        # Fourier transform we need to multiply g by a phase factor
        g = g * dr * np.exp(1j*k[:,None]*rmax) / (2*np.pi)**0.5
        return k, g

    # def twohalo(self, cosmo, bias, func, R, logm=np.logspace(12,16,41),
    #             bias_norm=1, **kwargs):
    def _twohalo(self, cosmo, offset_func, R, logm_2h=np.logspace(12,16,41),
                 z_2h=np.linspace(0,2,21), **kwargs):
        """Calculate the two-halo term associated with the profile

        Parameters
        ----------
        cosmo : `pyccl.Cosmology` object
        offset_func : callable or str
            if callable, it must be the offset version of the
            function in question. For instance, if one is
            modeling the convergence, the function supplied
            must be the offset convergence.
        R : np.ndarray
        logm : np.ndarray, optional
            masses over which to calculate the 2h term. If not
            specified, the masses used when defining the profile will be used.

        Notes
        -----
        kwargs are passed to the function to be called. If the function
        to be calculated is the convergence, the source redshift must
        be supplied
        """
        if not has_ccl:
            msg = 'Core Cosmology Library (CCL) required for two halo' \
                  ' calculations'
            raise ModuleNotFoundError(msg)
        assert isinstance(cosmo, ccl.Cosmology)
        # assert isinstance(bias, ccl.halos.HaloBias)
        # which function are we using?
        _valid_funcs = ('esd', 'kappa', 'sigma', 'convergence',
                        'projected', 'projected_cumulative',
                        'projected_excess',
                        'enclosed_surface_density', 'excess_surface_density',
                        'surface_density')
        if isinstance(offset_func, str):
            assert offset_func in _valid_funcs, \
                f'offset_func must be one of {_valid_funcs}'
            if offset_func in ('sigma', 'projected', 'surface_density'):
                offset_func = self.offset_projected
            elif offset_func in ('projected_cumulative',
                                 'enclosed_surface_density'):
                offset_func = self.offset_projected_cumulative
            elif offset_func in ('esd', 'projected_excess',
                                 'excess_surface_density'):
                offset_func = self.offset_projected_excess
            elif offset_func in ('kappa', 'convergence'):
                offset_func = self.offset_convergence
        else:
            assert callable(offset_func), \
                'argument offset_func must be a string or callable'
            assert offset_func.__name__.startswith('offset'), \
                'if argument offset_func is a function as opposed to' \
                'a string, it must be the offset version of the' \
                'function of interest.'
        # for each distance Ri, calculate the contribution from
        # all halos that will have some contribution at that distance
        # In order to speed this up a little, consider I don't need
        # to calculate the offset profile for all halos at all distances
        # but only for those who are close enough to that distance
        # So probably what I actually need to do is just calculate
        # for a grid in Roff and then move that to the required distance.
        # For example,
        # for i, Ri in R:
        #     twoh = offset_func(R, R[:i+1], **kwargs)
        return

    ### aliases for backward compatibility

    def surface_density(self, *args, **kwargs):
        """Alias for ``self.projected``"""
        return self.projected(*args, **kwargs)

    def enclosed_surface_density(self, *args, **kwargs):
        """Alias for ``self.projected_cumulative``"""
        return self.projected_cumulative(*args, **kwargs)

    def excess_surface_density(self, *args, **kwargs):
        """Alias for ``self.projected_excess``"""
        return self.projected_excess(*args, **kwargs)

    def offset_surface_density(self, *args, **kwargs):
        """Alias for ``self.offset_projected``"""
        return self.offset_projected(*args, **kwargs)

    def offset_enclosed_surface_density(self, *args, **kwargs):
        """Alias for ``self.offset_projected_cumulative``"""
        return self.offset_projected_cumulative(*args, **kwargs)

    def offset_excess_surface_density(self, *args, **kwargs):
        """Alias for ``self.offset_projected_excess``"""
        return self.offset_projected_excess(*args, **kwargs)

    ### auxiliary methods to test integration performance ###

    @inMpc
    @array
    def _quad_projected(self, R):
        """Not yet implemented"""
        integrand = lambda r, Ro: self.profile((r**2+Ro**2)**0.5)
        return np.array([[quad(integrand, 0, np.inf, args=(Rij,))
                          for Rij in Ri] for Ri in R])

    @inMpc
    @array
    def _test_integration(self, R, output=None):
        """Test the fast-integration methods against the slow
        but accurate quad function

        Not yet implemented
        """
        qsd = self.quad_projected(R)
        sd = self.projected(R)
        return

max_stars_repo_name: cristobal-sifon/profiley

max_issues_repo_name: cristobal-sifon/profiley

max_forks_repo_name: cristobal-sifon/profiley

reason: import numpy,from scipy,from astropy


text:
import jax
import jax.numpy as jnp

def _is_scalar(x):
  return jax.numpy.isscalar(x) or x.ndim == 0

class SDEProblem:
  '''
  Stores a vector stochastic differential equation of the form:
  :math:`d\\mathbf{x} = \\mathbf{a(x)} dt + \\mathbf{B(x)} d\\mathbf{w}`, where
  ``a`` and ``B`` are vector of drifts and matrix of noise coefficients.

  ``a`` vector and ``B`` matrix should be formatted such that

  .. prompt:: python >>> auto

    >>> x = x + a(x)*dt + np.dot(B(x),dw)

  is a valid code fragment (note ``B`` has as many columns as there are noise terms).

  Parameters
  ----------
  a : callable
    Function describing drift term of the equation, should return np.array of length ``dimension``.
  b : callable
      Function describing noise term of the equation, should return np.array of size ``(dimension,noiseterms)``.
  dimension : int
      Dimension of the space in which :math:`d\\mathbf{x}` sits.
  noiseterms : int
      Dimension of the space in which :math:`d\\mathbf{w}` sits.
  x0 : np.array
      Initial value of the stochastic process.
  tmax : float
      Time at which integration should stop.

  Example
  -------
  >>> import numpy as np
  >>> problem = pychastic.sde_problem.VectorSDEProblem(lambda x: np.array([1,1]), lambda x: np.array([[1,0.5],[0.5,1]]), 2, 2, np.array([1.5,0.5]), 1)

  '''
  def __init__(self, a, b, x0, tmax, exact_solution=None):

    if not tmax > 0:
      raise ValueError(f'tmax has to bo posiitve, not {tmax}')
    self.tmax = tmax

    # dimension & shape validation

    x0 = jnp.array(x0, dtype=jax.numpy.float32)

    (a_x0_shape, a_x0_dtype) = (jax.eval_shape(a,x0).shape,jax.eval_shape(a,x0).dtype)
    (b_x0_shape, b_x0_dtype) = (jax.eval_shape(b,x0).shape,jax.eval_shape(b,x0).dtype)

    if x0.ndim == len(a_x0_shape) == len(b_x0_shape) == 0:
      self.x0 = x0.reshape(1)
      # scalar case
      if len(a_x0_shape) != 1:
        new_a = lambda x: a(x).reshape(1)
        # TODO warn about reshaping
      if len(b_x0_shape) != 2:
        new_b = lambda x: b(x).reshape(1, 1)
        # TODO warn about reshaping

      self.dimension = self.noise_terms = 1
      self.a = new_a
      self.b = new_b

    elif x0.ndim == len(a_x0_shape) == 1 and len(b_x0_shape) == 2:
      # vector case
      if not x0.shape[0] == a_x0_shape[0] == b_x0_shape[0]:
        raise ValueError(f'Incosistent shapes: {x0.shape}, {a_x0_shape}, {b_x0_shape}')

      self.x0 = x0
      self.a = a
      self.b = b
      self.dimension, self.noise_terms = b_x0_shape

    else:
      raise ValueError(f'Inconsistent dimensions: {x0.shape}, {a_x0_shape}, {b_x0_shape}')

    if not isinstance(x0, jnp.ndarray):
        raise ValueError(f"{x0} should be jnp.array, not {type(x0)}")

    # dtype validation
    for val, key in [(x0.dtype, 'initial conditon'), (a_x0_dtype, 'drift term'), (b_x0_dtype, 'noise term')]:
      if not jnp.issubdtype(x0.dtype, jnp.floating):
        raise ValueError(f"{key} dtype should be float, not {val.dtype}.")

    # exact solution validation
    # TODO
    self.exact_solution = exact_solution

SDEProblem(lambda x: 1, lambda x: 2, 3, 4)

max_stars_repo_name: RadostW/stochastic

max_issues_repo_name: RadostW/stochastic

max_forks_repo_name: RadostW/stochastic

reason: import numpy,import jax


text:
"""
Generate a 3x3 heatmap where each square is the average (over N runs) # of live cancer cells.
The x-axis has the 3 migration bias values (0.25, 0.5, 0.75) and
the y-axis has the 3 attachment rate values (0.033, 0.2, 1).

Each of the 3x3 squares is annotated with N=# of runs (random seeds) that completed.

Input argument:
  fixed attachment lifetime (use the middle value = 60 for the heatmap in the BMC paper).

Output
  a .png at 300 dpi.

Module dependencies: numpy, matplotlib

Author:
Randy Heiland
Intelligent Systems Engineering, Indiana University

"""

import sys
import matplotlib.pyplot as plt
import matplotlib as mpl
from matplotlib.colors import BoundaryNorm
from matplotlib.ticker import MaxNLocator
from matplotlib.patches import Rectangle
import numpy as np

mpl.rcParams['savefig.dpi'] = 300

print("len(sys.argv)=", len(sys.argv))
if len(sys.argv) < 2:
  print("Usage: %s <attachment lifetime: 15, 60, or 120>" % sys.argv[0])
  sys.exit(0)

attachment_lifetime = float(sys.argv[1])
print("attachment_lifetime=", attachment_lifetime)

# -------- read in metrics file:
v = np.loadtxt("final_analysis.txt", delimiter=',')
run_num = v[0]
num_live_cancer_cells = v[1]
num_live_cancer_cells_above_threshold = v[2]
mean_oncoprotein = v[3]

# -------- read in params file:
#  -- on each line, we have 4 values: run#, rate, lifetime, bias
#     where:
# Attachment_rate = [0.033, 0.2, 1] - mean wait time, hours
# Attachment_lifetime = [15, 60, 120] - mins
# Migration_bias = [0.25, 0.5, 0.75] - random migration bias towards chemotactic gradients

run, rate, lifetime, bias = np.genfromtxt('run_rate_lifetime_bias.txt', delimiter=',').transpose()

# compute average # of <metric values>, e.g. # live cancer cells, mean oncoprotein,
m1 = np.zeros([3,3])
N_val = np.zeros([3,3])

for dummy in [1]:  # each bias value generates entire set of 3 heatmaps (for 3 metrics)
  iy = 0
  for rval in [0.033, 0.2, 1]:  # from bottom-to-top (Y-axis) - attachment rate
    ix = 0
    for bval in [0.25, 0.5, 0.75]:  # from left-to-rightop (X-axis) - migration bias
      idx = np.where( (lifetime==attachment_lifetime) & (rate==rval) & (bias==bval) )
      metric_vals = num_live_cancer_cells[idx]   # NB! change RHS for other metrics
      N = len(metric_vals)  # length of array = # of (random seed) runs successful
      N_val[iy][ix] = N
      avg_metric_vals = metric_vals.sum() / N
      print("N, avg_metric_vals = ",N, avg_metric_vals)
      m1[iy][ix] = avg_metric_vals
      ix += 1
    iy += 1

#-------------------------------------------------
#---- plotting: coarse-grained (3x3) mesh, color-mapped for "Avg(# live cells) over N runs" (N=max 10 = random seed range)
# make these smaller to increase the resolution
dx, dy = 1.0, 1.0

# generate 2 2d grids for the x & y bounds
y, x = np.mgrid[slice(0, 3 + dy, dy),
                slice(0, 3 + dx, dx)]

# x and y are bounds, so m1 is the value *inside* those bounds.
levels = MaxNLocator(nbins=128).tick_values(m1.min(), m1.max())

# pick the desired colormap, sensible levels, and define a normalization
# instance which takes data values and translates those into levels.
#cmap = plt.get_cmap('PiYG')
cmap = plt.get_cmap('plasma')
norm = BoundaryNorm(levels, ncolors=cmap.N, clip=True)

fig, ax = plt.subplots(nrows=1)

im = ax.pcolormesh(x, y, m1, cmap=cmap, norm=norm)
fig.colorbar(im, ax=ax)
title_string = "Avg(# live cancer cells) per N runs; Attachment Lifetime=%d" % (attachment_lifetime)
ax.set_title(title_string)
ax.set_xlabel("migration bias")
ax.set_ylabel("attachment rate")

#ax.set_xticks([])
#ax.set_yticks([])
plt.xticks([0.5, 1.5, 2.5], ["0.25", "0.5", "0.75"])
plt.yticks([0.5, 1.5, 2.5], ["0.033", "0.2", "1"])

#-----------
t = ax.transData
canvas = ax.figure.canvas
for iy in range(3):
  for ix in range(3):
    sval = "N=%d" % N_val[iy][ix]
    text = ax.text(ix+0.4, iy+0.1, sval, backgroundcolor="white", color="black", transform=t)
    text.set_bbox(dict(facecolor='white'))
#    text.set_bbox(dict(boxstyle='round,pad=1'))
    text.draw(canvas.get_renderer())

#text = ax.text(1.4, 0.1, "N=9", color="white", transform=t)
#text.draw(canvas.get_renderer())
#text = ax.text(2.4, 0.1, "N=8", color="white", transform=t)
#text.draw(canvas.get_renderer())

# adjust spacing between subplots so `ax1` title and `ax0` tick labels
# don't overlap
fig.tight_layout()

png_filename = "heatmap_avg_num_live_cancer_cells_lifetime%d.png" % (attachment_lifetime)
print("---> ", png_filename)
fig.savefig(png_filename, bbox_inches='tight')

#plt.show()

reason: import numpy


text:
import sys

import numpy as np
from skimage import measure

sys.path.append("../")

def test_clustering_widget(make_napari_viewer):

    import napari_clusters_plotter as ncp

    viewer = make_napari_viewer(strict_qt=True)
    widget_list = ncp.napari_experimental_provide_dock_widget()
    n_wdgts = len(viewer.window._dock_widgets)

    for widget in widget_list:
        _widget = widget(viewer)

        if isinstance(
            _widget, ncp._dimensionality_reduction.DimensionalityReductionWidget
        ):
            plot_widget = _widget

    viewer.window.add_dock_widget(plot_widget)
    assert len(viewer.window._dock_widgets) == n_wdgts + 1

def test_bad_measurements(qtbot, make_napari_viewer):

    from napari_clusters_plotter._dimensionality_reduction import (
        DimensionalityReductionWidget,
    )
    from napari_clusters_plotter._utilities import set_features

    label = np.array(
        [
            [0, 0, 0, 0, 0, 0, 0],
            [0, 1, 1, 0, 0, 2, 2],
            [0, 0, 0, 0, 2, 2, 2],
            [3, 3, 0, 0, 0, 0, 0],
            [0, 0, 4, 4, 0, 0, 0],
            [6, 6, 6, 6, 0, 5, 0],  # <-single pixel label leading to NaN meas.
            [0, 7, 7, 0, 0, 0, 0],
        ]
    )

    viewer = make_napari_viewer(strict_qt=True)
    labels_layer = viewer.add_labels(label)

    # Add NaNs to data
    measurements = measure.regionprops_table(
        label, properties=(["label", "area", "perimeter"])
    )
    for key in list(measurements.keys())[1:]:
        measurements[key] = measurements[key].astype(float)
        measurements[key][4] = np.nan
    set_features(labels_layer, measurements)

    widget = DimensionalityReductionWidget(napari_viewer=viewer)
    widget.run(
        viewer=viewer,
        labels_layer=labels_layer,
        selected_measurements_list=list(measurements.keys()),
        n_neighbours=2,
        perplexity=5,
        selected_algorithm="UMAP",
        standardize=False,
        n_components=2,
        explained_variance=95.0,
        pca_components=0,
    )

    blocker = qtbot.waitSignal(widget.worker.finished, timeout=1000000)
    blocker.wait()

# def test_call_to_function(qtbot, make_napari_viewer):

#     viewer = make_napari_viewer(strict_qt=True)

#     label = np.array(
#         [
#             [0, 0, 0, 0, 0, 0, 0],
#             [0, 1, 1, 0, 0, 2, 2],
#             [0, 0, 0, 0, 2, 2, 2],
#             [3, 3, 0, 0, 0, 0, 0],
#             [0, 0, 4, 4, 0, 5, 5],
#             [6, 6, 6, 6, 0, 5, 0],
#             [0, 7, 7, 0, 0, 0, 0],
#         ]
#     )

#     props = measure.regionprops_table(
#         label, properties=(["label", "area", "perimeter"])
#     )
#     label_layer = viewer.add_labels(label, properties=props)

#     from napari_clusters_plotter._dimensionality_reduction import (
#         DimensionalityReductionWidget,
#     )
#     from napari_clusters_plotter._utilities import get_layer_tabular_data

#     widget = DimensionalityReductionWidget(napari_viewer=viewer)
#     widget.run(
#         viewer=viewer,
#         labels_layer=label_layer,
#         selected_measurements_list=["area", "perimeter"],
#         n_neighbours=2,
#         perplexity=5,
#         selected_algorithm="UMAP",
#         standardize=False,
#         n_components=2,
#         explained_variance=95.0,
#         pca_components=0,
#     )

#     # waiting till the thread worker finished
#     blocker = qtbot.waitSignal(widget.worker.finished, timeout=1000000)
#     blocker.wait()
#     # additional waiting so the return_func_umap gets the returned embedding
#     # from the thread, and writes the results into properties/features of the labels layer
#     time.sleep(5)
#     result = get_layer_tabular_data(label_layer)

#     assert "UMAP_0" in result.columns
#     assert "UMAP_1" in result.columns

#     widget.run(
#         viewer=viewer,
#         labels_layer=label_layer,
#         selected_measurements_list=["area", "perimeter"],
#         n_neighbours=2,
#         perplexity=5,
#         selected_algorithm="t-SNE",
#         standardize=False,
#         n_components=2,
#         explained_variance=95.0,
#         pca_components=0,
#     )

#     blocker = qtbot.waitSignal(widget.worker.finished, timeout=1000000)
#     blocker.wait()
#     time.sleep(5)

#     result = get_layer_tabular_data(label_layer)
#     assert "t-SNE_0" in result.columns
#     assert "t-SNE_1" in result.columns

#     widget.run(
#         viewer=viewer,
#         labels_layer=label_layer,
#         selected_measurements_list=["area", "perimeter"],
#         n_neighbours=2,
#         perplexity=5,
#         selected_algorithm="PCA",
#         standardize=False,
#         n_components=2,
#         explained_variance=95.0,
#         pca_components=2,
#     )

#     blocker = qtbot.waitSignal(widget.worker.finished, timeout=10000000)
#     blocker.wait()
#     time.sleep(10)

#     result = get_layer_tabular_data(label_layer)
#     assert "PC_0" in result.columns

def test_umap():

    import pandas as pd

    from napari_clusters_plotter._dimensionality_reduction import umap

    X = np.array([[0, 0, 0], [0, 1, 1], [1, 0, 1], [1, 1, 1]])
    n_comp = 2

    # umap returns (str, np.ndarray), where the first item is algorithm name
    result = umap(pd.DataFrame(X), n_neigh=2, n_components=n_comp)

    assert result[1].shape[-1] == n_comp

def test_tsne():

    X = np.array([[0, 0, 0], [0, 1, 1], [1, 0, 1], [1, 1, 1]])
    n_comp = 2

    import pandas as pd

    from napari_clusters_plotter._dimensionality_reduction import tsne

    result = tsne(pd.DataFrame(X), perplexity=5, n_components=2)
    assert result[1].shape[-1] == n_comp

def test_pca():

    X = np.array([[0, 0, 0], [0, 1, 1], [1, 0, 1], [1, 1, 1]])
    n_comp = 3

    import pandas as pd

    from napari_clusters_plotter._dimensionality_reduction import pca

    result = pca(pd.DataFrame(X), explained_variance_threshold=95.0, n_components=0)
    assert result[1].shape[-1] == n_comp

    result = pca(pd.DataFrame(X), explained_variance_threshold=95.0, n_components=0)
    assert result[1].shape[-1] == n_comp

if __name__ == "__main__":
    pass

    # test_clustering_widget()
    import napari

    test_bad_measurements(napari.Viewer)
    # test_call_to_function(napari.Viewer)
    # test_umap()

max_stars_repo_name: BiAPoL/napari-clusters-plotter

max_issues_repo_name: BiAPoL/napari-clusters-plotter

max_forks_repo_name: BiAPoL/napari-clusters-plotter