pimco-charts/app/renderer/scale.py

"""Data-to-pixel coordinate mapping."""

from __future__ import annotations
import math
from datetime import datetime
from typing import Sequence


class LinearScale:
    """Maps a numeric data domain to a pixel range."""

    def __init__(self, domain_min: float, domain_max: float, range_min: float, range_max: float):
        self.domain_min = domain_min
        self.domain_max = domain_max
        self.range_min = range_min
        self.range_max = range_max
        span = domain_max - domain_min
        self._scale = (range_max - range_min) / span if span != 0 else 0

    def __call__(self, value: float) -> float:
        return self.range_min + (value - self.domain_min) * self._scale

    def invert(self, pixel: float) -> float:
        if self._scale == 0:
            return self.domain_min
        return self.domain_min + (pixel - self.range_min) / self._scale


class DateScale:
    """Maps datetime values to pixel positions."""

    def __init__(self, date_min: datetime, date_max: datetime, range_min: float, range_max: float):
        self.date_min = date_min
        self.date_max = date_max
        self.range_min = range_min
        self.range_max = range_max
        total_seconds = (date_max - date_min).total_seconds()
        self._scale = (range_max - range_min) / total_seconds if total_seconds != 0 else 0

    def __call__(self, value: datetime) -> float:
        seconds = (value - self.date_min).total_seconds()
        return self.range_min + seconds * self._scale

    def invert(self, pixel: float) -> datetime:
        from datetime import timedelta
        if self._scale == 0:
            return self.date_min
        seconds = (pixel - self.range_min) / self._scale
        return self.date_min + timedelta(seconds=seconds)


class CategoricalScale:
    """Maps category labels (strings) to evenly spaced pixel centre positions."""

    def __init__(self, categories: list[str], range_min: float, range_max: float):
        self.categories = list(categories)
        self.range_min = range_min
        self.range_max = range_max
        n = len(self.categories)
        self._step = (range_max - range_min) / n if n > 0 else 0
        self._positions = {
            cat: range_min + self._step * (i + 0.5)
            for i, cat in enumerate(self.categories)
        }

    def __call__(self, value: str) -> float:
        return self._positions.get(str(value), self.range_min)

    @property
    def step(self) -> float:
        return self._step


def nice_ticks(data_min: float, data_max: float, target_count: int = 6) -> list[float]:
    """Generate aesthetically pleasing tick values for an axis."""
    if data_min == data_max:
        return [data_min]

    raw_step = (data_max - data_min) / target_count
    magnitude = 10 ** math.floor(math.log10(abs(raw_step)))
    residual = raw_step / magnitude

    if residual <= 1.5:
        nice_step = 1 * magnitude
    elif residual <= 3:
        nice_step = 2 * magnitude
    elif residual <= 7:
        nice_step = 5 * magnitude
    else:
        nice_step = 10 * magnitude

    tick_min = math.floor(data_min / nice_step) * nice_step
    tick_max = math.ceil(data_max / nice_step) * nice_step

    ticks = []
    val = tick_min
    while val <= tick_max + nice_step * 0.01:
        ticks.append(round(val, 10))
        val += nice_step

    return ticks


def nice_date_ticks(date_min: datetime, date_max: datetime, target_count: int = 8) -> list[datetime]:
    """Generate evenly spaced date ticks."""
    total_days = (date_max - date_min).days
    if total_days <= 0:
        return [date_min]

    from datetime import timedelta
    from dateutil.relativedelta import relativedelta

    # Choose interval based on span
    if total_days <= 365:
        # Monthly ticks
        interval = relativedelta(months=1)
    elif total_days <= 365 * 3:
        # Quarterly
        interval = relativedelta(months=3)
    elif total_days <= 365 * 8:
        # Semi-annual
        interval = relativedelta(months=6)
    elif total_days <= 365 * 20:
        # Annual
        interval = relativedelta(years=1)
    else:
        # Every 2 years
        interval = relativedelta(years=2)

    # Start from the first day of the month at or after date_min
    current = datetime(date_min.year, date_min.month, 1)
    if current < date_min:
        current += relativedelta(months=1)

    ticks = []
    while current <= date_max:
        ticks.append(current)
        current += interval

    # Thin out if too many ticks
    while len(ticks) > target_count * 1.5:
        ticks = ticks[::2]

    return ticks