Re-org GeMOSA codes

xautodl/datasets/math_core.py

@@ -4,6 +4,7 @@
 from .math_base_funcs import LinearFunc, QuadraticFunc, CubicFunc, QuarticFunc
 from .math_dynamic_funcs import DynamicLinearFunc
 from .math_dynamic_funcs import DynamicQuadraticFunc
+from .math_dynamic_funcs import DynamicSinQuadraticFunc
 from .math_adv_funcs import ConstantFunc
 from .math_adv_funcs import ComposedSinFunc, ComposedCosFunc
 from .math_dynamic_generator import GaussianDGenerator

xautodl/datasets/math_dynamic_funcs.py

@@ -5,9 +5,6 @@ import math
 import abc
 import copy
 import numpy as np
-from typing import Optional
-import torch
-import torch.utils.data as data
 
 from .math_base_funcs import FitFunc
 
@@ -68,10 +65,11 @@ class DynamicQuadraticFunc(DynamicFunc):
     def __init__(self, params=None):
         super(DynamicQuadraticFunc, self).__init__(3, params)
 
-    def __call__(self, x, timestamp=None):
+    def __call__(
+        self,
+        x,
+    ):
         self.check_valid()
-        if timestamp is None:
-            timestamp = self._timestamp
         a = self._params[0](timestamp)
         b = self._params[1](timestamp)
         c = self._params[2](timestamp)
@@ -80,10 +78,38 @@ class DynamicQuadraticFunc(DynamicFunc):
         return a * x * x + b * x + c
 
     def __repr__(self):
-        return "{name}({a} * x^2 + {b} * x + {c}, timestamp={timestamp})".format(
+        return "{name}({a} * x^2 + {b} * x + {c})".format(
+            name=self.__class__.__name__,
+            a=self._params[0],
+            b=self._params[1],
+            c=self._params[2],
+        )
+
+
+class DynamicSinQuadraticFunc(DynamicFunc):
+    """The dynamic quadratic function that outputs f(x) = sin(a * x^2 + b * x + c).
+    The a, b, and c is a function of timestamp.
+    """
+
+    def __init__(self, params=None):
+        super(DynamicSinQuadraticFunc, self).__init__(3, params)
+
+    def __call__(
+        self,
+        x,
+    ):
+        self.check_valid()
+        a = self._params[0](timestamp)
+        b = self._params[1](timestamp)
+        c = self._params[2](timestamp)
+        convert_fn = lambda x: x[-1] if isinstance(x, (tuple, list)) else x
+        a, b, c = convert_fn(a), convert_fn(b), convert_fn(c)
+        return math.sin(a * x * x + b * x + c)
+
+    def __repr__(self):
+        return "{name}({a} * x^2 + {b} * x + {c})".format(
             name=self.__class__.__name__,
             a=self._params[0],
             b=self._params[1],
             c=self._params[2],
-            timestamp=self._timestamp,
         )

xautodl/datasets/synthetic_core.py

@@ -3,7 +3,7 @@ from .synthetic_utils import TimeStamp
 from .synthetic_env import SyntheticDEnv
 from .math_core import LinearFunc
 from .math_core import DynamicLinearFunc
-from .math_core import DynamicQuadraticFunc
+from .math_core import DynamicQuadraticFunc, DynamicSinQuadraticFunc
 from .math_core import (
     ConstantFunc,
     ComposedSinFunc as SinFunc,
@@ -63,9 +63,9 @@ def get_synthetic_env(total_timestamp=1600, num_per_task=1000, mode=None, versio
         time_generator = TimeStamp(
             min_timestamp=0, max_timestamp=max_time, num=total_timestamp, mode=mode
         )
-        oracle_map = DynamicQuadraticFunc(
+        oracle_map = DynamicSinQuadraticFunc(
             params={
-                0: LinearFunc(params={0: 0.1, 1: 0}),  # 0.1 * t
+                0: CosFunc(params={0: 0.5, 1: 1, 2: 1}),  # 0.5 cos(t) + 1
                 1: SinFunc(params={0: 1, 1: 1, 2: 0}),  # sin(t)
                 2: ConstantFunc(0),
             }

xautodl/datasets/synthetic_env.py

@@ -1,6 +1,3 @@
-import math
-import random
-from typing import List, Optional, Dict
 import torch
 import torch.utils.data as data
 
@@ -43,6 +40,18 @@ class SyntheticDEnv(data.Dataset):
         self._oracle_map = oracle_map
         self._num_per_task = num_per_task
         self._noise = noise
+        self._meta_info = dict()
+
+    def set_regression(self):
+        self._meta_info["task"] = "regression"
+
+    def set_classification(self, num_classes):
+        self._meta_info["task"] = "classification"
+        self._meta_info["num_classes"] = int(num_classes)
+
+    @property
+    def meta_info(self):
+        return self._meta_info
 
     @property
     def min_timestamp(self):
@@ -60,13 +69,6 @@ class SyntheticDEnv(data.Dataset):
     def mode(self):
         return self._time_generator.mode
 
-    def random_timestamp(self, min_timestamp=None, max_timestamp=None):
-        if min_timestamp is None:
-            min_timestamp = self.min_timestamp
-        if max_timestamp is None:
-            max_timestamp = self.max_timestamp
-        return random.random() * (max_timestamp - min_timestamp) + min_timestamp
-
     def get_timestamp(self, index):
         if index is None:
             timestamps = []