autodl-projects/lib/datasets/synthetic_adaptive_environment.py

#####################################################
# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2021.03 #
#####################################################
import math
import abc
import numpy as np
from typing import Optional
import torch
import torch.utils.data as data


class FitFunc(abc.ABC):
    """The fit function that outputs f(x) = a * x^2 + b * x + c."""

    def __init__(self, freedom: int, list_of_points=None):
        self._params = dict()
        for i in range(freedom):
            self._params[i] = None
        self._freedom = freedom
        if list_of_points is not None:
            self.fit(list_of_points)

    def set(self, _params):
        self._params = copy.deepcopy(_params)

    def check_valid(self):
        for key, value in self._params.items():
            if value is None:
                raise ValueError("The {:} is None".format(key))

    @abc.abstractmethod
    def __getitem__(self, x):
        raise NotImplementedError

    @abc.abstractmethod
    def _getitem(self, x):
        raise NotImplementedError

    def fit(
        self,
        list_of_points,
        max_iter=900,
        lr_max=1.0,
        verbose=False,
    ):
        with torch.no_grad():
            data = torch.Tensor(list_of_points).type(torch.float32)
            assert data.ndim == 2 and data.size(1) == 2, "Invalid shape : {:}".format(
                data.shape
            )
            x, y = data[:, 0], data[:, 1]
        weights = torch.nn.Parameter(torch.Tensor(self._freedom))
        torch.nn.init.normal_(weights, mean=0.0, std=1.0)
        optimizer = torch.optim.Adam([weights], lr=lr_max, amsgrad=True)
        lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
            optimizer,
            milestones=[
                int(max_iter * 0.25),
                int(max_iter * 0.5),
                int(max_iter * 0.75),
            ],
            gamma=0.1,
        )
        if verbose:
            print("The optimizer: {:}".format(optimizer))

        best_loss = None
        for _iter in range(max_iter):
            y_hat = self._getitem(x, weights)
            loss = torch.mean(torch.abs(y - y_hat))
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            lr_scheduler.step()
            if verbose:
                print(
                    "In the fit, loss at the {:02d}/{:02d}-th iter is {:}".format(
                        _iter, max_iter, loss.item()
                    )
                )
            # Update the params
            if best_loss is None or best_loss > loss.item():
                best_loss = loss.item()
                for i in range(self._freedom):
                    self._params[i] = weights[i].item()

    def __repr__(self):
        return "{name}(freedom={freedom})".format(
            name=self.__class__.__name__, freedom=freedom
        )


class QuadraticFunc(FitFunc):
    """The quadratic function that outputs f(x) = a * x^2 + b * x + c."""

    def __init__(self, list_of_points=None):
        super(QuadraticFunc, self).__init__(3, list_of_points)

    def __getitem__(self, x):
        self.check_valid()
        return self._params[0] * x * x + self._params[1] * x + self._params[2]

    def _getitem(self, x, weights):
        return weights[0] * x * x + weights[1] * x + weights[2]

    def __repr__(self):
        return "{name}(y = {a} * x^2 + {b} * x + {c})".format(
            name=self.__class__.__name__,
            a=self._params[0],
            b=self._params[1],
            c=self._params[2],
        )


class CubicFunc(FitFunc):
    """The cubic function that outputs f(x) = a * x^3 + b * x^2 + c * x + d."""

    def __init__(self, list_of_points=None):
        super(CubicFunc, self).__init__(4, list_of_points)

    def __getitem__(self, x):
        self.check_valid()
        return (
            self._params[0] * x ** 3
            + self._params[1] * x ** 2
            + self._params[2] * x
            + self._params[3]
        )

    def _getitem(self, x, weights):
        return weights[0] * x ** 3 + weights[1] * x ** 2 + weights[2] * x + weights[3]

    def __repr__(self):
        return "{name}(y = {a} * x^3 + {b} * x^2 + {c} * x + {d})".format(
            name=self.__class__.__name__,
            a=self._params[0],
            b=self._params[1],
            c=self._params[2],
            d=self._params[3],
        )


class QuarticFunc(FitFunc):
    """The quartic function that outputs f(x) = a * x^4 + b * x^3 + c * x^2 + d * x + e."""

    def __init__(self, list_of_points=None):
        super(QuarticFunc, self).__init__(5, list_of_points)

    def __getitem__(self, x):
        self.check_valid()
        return (
            self._params[0] * x ** 4
            + self._params[1] * x ** 3
            + self._params[2] * x ** 2
            + self._params[3] * x
            + self._params[4]
        )

    def _getitem(self, x, weights):
        return (
            weights[0] * x ** 4
            + weights[1] * x ** 3
            + weights[2] * x ** 2
            + weights[3] * x
            + weights[4]
        )

    def __repr__(self):
        return "{name}(y = {a} * x^4 + {b} * x^3 + {c} * x^2 + {d} * x + {e})".format(
            name=self.__class__.__name__,
            a=self._params[0],
            b=self._params[1],
            c=self._params[2],
            d=self._params[3],
            e=self._params[3],
        )


class SynAdaptiveEnv(data.Dataset):
    """The synethtic dataset for adaptive environment.

    - x in [0, 1]
    - y = amplitude-scale-of(x) * sin( period-phase-shift-of(x) )
    - where
    - the amplitude scale is a quadratic function of x
    - the period-phase-shift is another quadratic function of x

    """

    def __init__(
        self,
        num: int = 100,
        num_sin_phase: int = 7,
        min_amplitude: float = 1,
        max_amplitude: float = 4,
        phase_shift: float = 0,
        mode: Optional[str] = None,
    ):
        self._amplitude_scale = QuadraticFunc(
            [(0, min_amplitude), (0.5, max_amplitude), (1, min_amplitude)]
        )

        self._num_sin_phase = num_sin_phase
        self._interval = 1.0 / (float(num) - 1)
        self._total_num = num

        fitting_data = []
        temp_max_scalar = 2 ** (num_sin_phase - 1)
        for i in range(num_sin_phase):
            value = (2 ** i) / temp_max_scalar
            next_value = (2 ** (i + 1)) / temp_max_scalar
            for _phase in (0, 0.25, 0.5, 0.75):
                inter_value = value + (next_value - value) * _phase
                fitting_data.append((inter_value, math.pi * (2 * i + _phase)))
        self._period_phase_shift = QuarticFunc(fitting_data)

        # Training Set 60%
        num_of_train = int(self._total_num * 0.6)
        # Validation Set 20%
        num_of_valid = int(self._total_num * 0.2)
        # Test Set 20%
        num_of_set = self._total_num - num_of_train - num_of_valid
        all_indexes = list(range(self._total_num))
        if mode is None:
            self._indexes = all_indexes
        elif mode.lower() in ("train", "training"):
            self._indexes = all_indexes[:num_of_train]
        elif mode.lower() in ("valid", "validation"):
            self._indexes = all_indexes[num_of_train : num_of_train + num_of_valid]
        elif mode.lower() in ("test", "testing"):
            self._indexes = all_indexes[num_of_train + num_of_valid :]
        else:
            raise ValueError("Unkonwn mode of {:}".format(mode))

    def __iter__(self):
        self._iter_num = 0
        return self

    def __next__(self):
        if self._iter_num >= len(self):
            raise StopIteration
        self._iter_num += 1
        return self.__getitem__(self._iter_num - 1)

    def __getitem__(self, index):
        assert 0 <= index < len(self), "{:} is not in [0, {:})".format(index, len(self))
        index = self._indexes[index]
        position = self._interval * index
        value = self._amplitude_scale[position] * math.sin(
            self._period_phase_shift[position]
        )
        return index, position, value

    def __len__(self):
        return len(self._indexes)

    def __repr__(self):
        return (
            "{name}({cur_num:}/{total} elements,\n"
            "amplitude={amplitude},\n"
            "period_phase_shift={period_phase_shift})".format(
                name=self.__class__.__name__,
                cur_num=self._total_num,
                total=len(self),
                amplitude=self._amplitude_scale,
                period_phase_shift=self._period_phase_shift,
            )
        )