nmreval/nmreval/fit/parameter.py

from __future__ import annotations

from numbers import Number
from itertools import count

import numpy as np


class Parameters(dict):
    count = count()

    def __str__(self):
        return 'Parameters:\n' + '\n'.join([str(k)+': '+str(v) for k, v in self.items()])

    def __getitem__(self, item):
        if isinstance(item, (list, tuple, np.ndarray)):
            values = []
            for item_i in item:
                values.append(super().__getitem__(item_i))
            return values
        else:
            return super().__getitem__(item)

    @staticmethod
    def _prep_bounds(val, p_len: int) -> list:
        # helper function to ensure that bounds and variable are of parameter shape
        if isinstance(val, (Number, bool)) or val is None:
            return [val] * p_len

        elif len(val) == p_len:
            return val

        elif len(val) == 1:
            return [val[0]] * p_len

        else:
            raise ValueError('Input {} has wrong dimensions'.format(val))

    def add_parameter(self, param, var=None, lb=None, ub=None):
        if isinstance(param, Number):
            param = [param]

        p_len = len(param)

        # make list if only single value is given
        var = self._prep_bounds(var, p_len)
        lb = self._prep_bounds(lb, p_len)
        ub = self._prep_bounds(ub, p_len)

        new_keys = []
        for i in range(p_len):
            new_idx = next(self.count)
            new_keys.append(new_idx)

            self[new_idx] = Parameter(param[i], var=var[i], lb=lb[i], ub=ub[i])

        return new_keys

    def copy(self):
        p = Parameters()
        for k, v in self.items():
            p[k] = Parameter(v.value, var=v.var, lb=v.lb, ub=v.ub)

        if len(p) == 0:
            return p

        max_k = max(p.keys())
        c = next(p.count)
        while c < max_k:
            c = next(p.count)

        return p

    def get_state(self):
        return {k: v.get_state() for k, v in self.items()}


class Parameter:
    """
    Container for one parameter
    """
    __slots__ = ['name', 'value', 'error', 'init_val', 'var', 'lb', 'ub', 'scale', 'function']

    def __init__(self, value: float, var: bool = True, lb: float = -np.inf, ub: float = np.inf):
        self.lb = lb if lb is not None else -np.inf
        self.ub = ub if ub is not None else np.inf

        if self.lb <= value <= self.ub:
            self.value = value
        else:
            raise ValueError('Value of parameter is outside bounds')

        self.init_val = value

        with np.errstate(divide='ignore'):
            # throws RuntimeWarning for zeros
            self.scale = 10**(np.floor(np.log10(np.abs(self.value))))

        if self.scale == 0:
            self.scale = 1.

        self.var = bool(var) if var is not None else True
        self.error = None if self.var is False else 0.0
        self.name = ''
        self.function = ''

    def __str__(self):
        start = ''
        if self.name:
            if self.function:
                start = f'{self.name} ({self.function}):  '
            else:
                start = self.name + ':  '

        if self.var:
            return start + f'{self.value:.4g} +/- {self.error:.4g}, init={self.init_val}'
        else:
            return start + f'{self.value:} (fixed)'

    def __add__(self, other: Parameter | float) -> float:
        if isinstance(other, (float, int)):
            return self.value + other
        elif isinstance(other, Parameter):
            return self.value + other.value

    def __radd__(self, other: Parameter | float) -> float:
        return self.__add__(other)

    @property
    def scaled_value(self):
        return self.value / self.scale

    @scaled_value.setter
    def scaled_value(self, value):
        self.value = value * self.scale

    @property
    def scaled_error(self):
        if self.error is None:
            return self.error
        else:
            return self.error / self.scale

    @scaled_error.setter
    def scaled_error(self, value):
        self.error = value * self.scale

    def get_state(self):

        return {slot: getattr(self, slot) for slot in  self.__slots__}

    @staticmethod
    def set_state(state: dict):
        par = Parameter(state.pop('value'))
        for k, v in state.items():
            setattr(par, k, v)

        return par

    @property
    def full_name(self):
        name = self.name
        if self.function:
            name += ' (' + self.function + ')'

        return name