Source code for schrodinger.math.mathutils

"""
Contains math-related utility functions

Copyright Schrodinger, LLC. All rights reserved.
"""

import math
import numpy

from collections import defaultdict
from scipy import interpolate


def round_value(val, precision=3, significant_figures=None):
    """
    Return val as a string with the required precision or significant figures.

    Either precision or significant_figures should be provided.
    Uses scientific notation for very large or small values, if the precision allows.

    :param float val: The value to round

    :param int precision: The precision needed after rounding. A precision of 2 means
        two decimal places should be kept after rounding. A negative precision indicates
        how many digits can be replaced with zeros before the decimal point.
        -5 means that 123456 can be shown as 1e5, -3 means it can be shown as 1.23e5.

    :param int significant_figures: The number of significant figures that should
        remain after rounding. If provided, determines the rounding precision.

    :rtype: str or None
    :return: A string with the required precision, or None if the input is a string
    """

    if isinstance(val, str):
        return None

    val = float(val)
    if val == 0:
        return '0'

    magnitude = int(numpy.floor(numpy.log10(abs(val))))

    if significant_figures:
        precision = significant_figures - magnitude - 1

    # Uses scientific notation if more than 4 digits can be removed in large numbers
    if magnitude < -3 or (magnitude > 3 and precision < -3):
        decimals = max(0, magnitude + precision)
        return f'{val:.{decimals}e}'

    elif precision > 0:
        return f'{val:.{precision}f}'
    else:
        return str(round(val))


def deduplicate_xy_data(x_vals, y_vals):
    """
    Remove duplicate x values by averaging the y values for them.

    :param list x_vals: The x values

    :param list y_vals: The y values

    :rtype: list, list
    :return: The deduplicated xy data
    """
    all_vals = defaultdict(list)
    for x_val, y_val in zip(x_vals, y_vals):
        all_vals[x_val].append(y_val)
    deduped_x = list(all_vals.keys())
    deduped_y = [sum(all_vals[x]) / len(all_vals[x]) for x in deduped_x]
    return deduped_x, deduped_y


class Interpolate1D:
    """
    Creates a map between values in a source and a target axis, allowing to get
    the equivalent target point for each source point.

    Wrapper around `scipy.interpolate.interp1d` to allow logarithmic
    interpolation or extrapolation.
    """

    def __init__(self, source_vals, target_vals, log_interp=False):
        """
        Create an instance.

        :param tuple source_vals: The values of points in the source range

        :param tuple target_vals: The values of points in the target range

        :param bool log_interp: Whether the interpolation is logarithmic.
            If False, linear interpolation will be used.
        """
        if log_interp:
            target_vals = [numpy.log10(x) for x in target_vals]
            linear_interp = interpolate.interp1d(
                source_vals, target_vals, fill_value='extrapolate')
            self.interp = lambda x: numpy.power(10, linear_interp(x))
        else:
            self.interp = interpolate.interp1d(
                source_vals, target_vals, fill_value='extrapolate')

    def __call__(self, source_val):
        """
        Get the equivalent target value of the passed source value

        :param float source_val: The value in the source range

        :rtype: float
        :return: The equivalent value in the target range
        """
        return self.interp(source_val)


class Interpolate2D:
    """
    Creates two instances of Interpolate1D to map values between two source axes
    and two target axes. Example use case is mapping QGraphicsScene/QChart
    XY coordinates to a XY coordinate system being displayed in the scene/chart.

    The two axes need to be independent.
    """

    def __init__(self,
                 x_source_vals,
                 x_target_vals,
                 y_source_vals,
                 y_target_vals,
                 x_log_interp=False,
                 y_log_interp=False):
        """
        Create an instance.

        :param tuple x_source_vals: The values of points in the X source range

        :param tuple x_target_vals: The values of points in the X target range

        :param tuple y_source_vals: The values of points in the Y source range

        :param tuple y_target_vals: The values of points in the Y target range

        :param bool x_log_interp: Whether the X axis interpolation is logarithmic

        :param bool y_log_interp: Whether the Y axis interpolation is logarithmic
        """
        self.x_interp = Interpolate1D(x_source_vals, x_target_vals,
                                      x_log_interp)
        self.y_interp = Interpolate1D(y_source_vals, y_target_vals,
                                      y_log_interp)

    def __call__(self, source_x_val, source_y_val):
        """
        Get the equivalent target values of the passed source values

        :param float source_x_val: The X value in the source range

        :param float source_y_val: The Y value in the source range

        :rtype: float, float
        :return: The equivalent x, y values in the target range
        """
        return (self.x_interp(source_x_val), self.y_interp(source_y_val))