Update the sync data v1

xautodl/datasets/synthetic_env.py

@@ -1,15 +1,9 @@
-#####################################################
-# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2021.04 #
-#####################################################
 import math
 import random
-import numpy as np
 from typing import List, Optional, Dict
 import torch
 import torch.utils.data as data
 
-from .synthetic_utils import TimeStamp
-
 
 def is_list_tuple(x):
     return isinstance(x, (tuple, list))
@@ -38,46 +32,33 @@ class SyntheticDEnv(data.Dataset):
 
     def __init__(
         self,
-        mean_functors: List[data.Dataset],
-        cov_functors: List[List[data.Dataset]],
+        data_generator,
+        oracle_map,
+        time_generator,
         num_per_task: int = 5000,
-        timestamp_config: Optional[Dict] = None,
-        mode: Optional[str] = None,
-        timestamp_noise_scale: float = 0.3,
+        noise: float = 0.1,
     ):
-        self._ndim = len(mean_functors)
-        assert self._ndim == len(
-            cov_functors
-        ), "length does not match {:} vs. {:}".format(self._ndim, len(cov_functors))
-        for cov_functor in cov_functors:
-            assert self._ndim == len(
-                cov_functor
-            ), "length does not match {:} vs. {:}".format(self._ndim, len(cov_functor))
+        self._data_generator = data_generator
+        self._time_generator = time_generator
+        self._oracle_map = oracle_map
         self._num_per_task = num_per_task
-        if timestamp_config is None:
-            timestamp_config = dict(mode=mode)
-        elif "mode" not in timestamp_config:
-            timestamp_config["mode"] = mode
-
-        self._timestamp_generator = TimeStamp(**timestamp_config)
-        self._timestamp_noise_scale = timestamp_noise_scale
-
-        self._mean_functors = mean_functors
-        self._cov_functors = cov_functors
-
-        self._oracle_map = None
+        self._noise = noise
 
     @property
     def min_timestamp(self):
-        return self._timestamp_generator.min_timestamp
+        return self._time_generator.min_timestamp
 
     @property
     def max_timestamp(self):
-        return self._timestamp_generator.max_timestamp
+        return self._time_generator.max_timestamp
 
     @property
-    def timestamp_interval(self):
-        return self._timestamp_generator.interval
+    def time_interval(self):
+        return self._time_generator.interval
+
+    @property
+    def mode(self):
+        return self._time_generator.mode
 
     def random_timestamp(self, min_timestamp=None, max_timestamp=None):
         if min_timestamp is None:
@@ -89,16 +70,13 @@ class SyntheticDEnv(data.Dataset):
     def get_timestamp(self, index):
         if index is None:
             timestamps = []
-            for index in range(len(self._timestamp_generator)):
-                timestamps.append(self._timestamp_generator[index][1])
+            for index in range(len(self._time_generator)):
+                timestamps.append(self._time_generator[index][1])
             return tuple(timestamps)
         else:
-            index, timestamp = self._timestamp_generator[index]
+            index, timestamp = self._time_generator[index]
             return timestamp
 
-    def set_oracle_map(self, functor):
-        self._oracle_map = functor
-
     def __iter__(self):
         self._iter_num = 0
         return self
@@ -111,7 +89,7 @@ class SyntheticDEnv(data.Dataset):
 
     def __getitem__(self, index):
         assert 0 <= index < len(self), "{:} is not in [0, {:})".format(index, len(self))
-        index, timestamp = self._timestamp_generator[index]
+        index, timestamp = self._time_generator[index]
         return self.__call__(timestamp)
 
     def seq_call(self, timestamps):
@@ -122,52 +100,24 @@ class SyntheticDEnv(data.Dataset):
             return zip_sequence(xdata)
 
     def __call__(self, timestamp):
-        mean_list = [functor(timestamp) for functor in self._mean_functors]
-        cov_matrix = [
-            [abs(cov_gen(timestamp)) for cov_gen in cov_functor]
-            for cov_functor in self._cov_functors
-        ]
-
-        dataset = np.random.multivariate_normal(
-            mean_list, cov_matrix, size=self._num_per_task
+        dataset = self._data_generator(timestamp, self._num_per_task)
+        targets = self._oracle_map.noise_call(dataset, timestamp, self._noise)
+        return torch.Tensor([timestamp]), (
+            torch.Tensor(dataset),
+            torch.Tensor(targets),
         )
-        if self._oracle_map is None:
-            return torch.Tensor([timestamp]), torch.Tensor(dataset)
-        else:
-            targets = self._oracle_map.noise_call(dataset, timestamp)
-            return torch.Tensor([timestamp]), (
-                torch.Tensor(dataset),
-                torch.Tensor(targets),
-            )
 
     def __len__(self):
-        return len(self._timestamp_generator)
+        return len(self._time_generator)
 
     def __repr__(self):
         return "{name}({cur_num:}/{total} elements, ndim={ndim}, num_per_task={num_per_task}, range=[{xrange_min:.5f}~{xrange_max:.5f}], mode={mode})".format(
             name=self.__class__.__name__,
             cur_num=len(self),
-            total=len(self._timestamp_generator),
+            total=len(self._time_generator),
             ndim=self._ndim,
             num_per_task=self._num_per_task,
             xrange_min=self.min_timestamp,
             xrange_max=self.max_timestamp,
-            mode=self._timestamp_generator.mode,
+            mode=self.mode,
         )
-
-
-class EnvSampler:
-    def __init__(self, env, batch, enlarge):
-        indexes = list(range(len(env)))
-        self._indexes = indexes * enlarge
-        self._batch = batch
-        self._iterations = len(self._indexes) // self._batch
-
-    def __iter__(self):
-        random.shuffle(self._indexes)
-        for it in range(self._iterations):
-            indexes = self._indexes[it * self._batch : (it + 1) * self._batch]
-            yield indexes
-
-    def __len__(self):
-        return self._iterations