Upgrade spaces and add more tests

.github/workflows/test.yml

@@ -24,6 +24,16 @@ jobs:
         with:
           python-version: ${{ matrix.python-version }}
 
+      - name: Lint with Black
+        run: |
+          python -m pip install black
+          python --version
+          python -m black --version
+          echo $PWD ; ls
+          python -m black ./tests -l 88 --check --diff --verbose
+          python -m black ./lib/spaces -l 88 --check --diff --verbose
+          python -m black ./lib/trade_models -l 88 --check --diff --verbose
+
       - name: Test Search Space
         run: |
           python -m pip install pytest

README.md

@@ -154,7 +154,7 @@ If you want to contribute to this repo, please see [CONTRIBUTING.md](.github/CON
 Besides, please follow [CODE-OF-CONDUCT.md](.github/CODE-OF-CONDUCT.md).
 
 We use [`black`](https://github.com/psf/black) for Python code formatter.
-Please use `black . -l 120`.
+Please use `black . -l 88`.
 
 # License
 The entire codebase is under the [MIT license](LICENSE.md).

lib/layers/super_mlp.py

@@ -1,7 +1,43 @@
 import torch.nn as nn
+from torch.nn.parameter import Parameter
 from typing import Optional
 
-class MLP(nn.Module):
+
+class Linear(nn.Module):
+    """Applies a linear transformation to the incoming data: :math:`y = xA^T + b`
+    """
+    __constants__ = ['in_features', 'out_features']
+    in_features: int
+    out_features: int
+    weight: Tensor
+
+    def __init__(self, in_features: int, out_features: int, bias: bool = True) -> None:
+        super(Linear, self).__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.weight = Parameter(torch.Tensor(out_features, in_features))
+        if bias:
+            self.bias = Parameter(torch.Tensor(out_features))
+        else:
+            self.register_parameter('bias', None)
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+        if self.bias is not None:
+            fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
+            bound = 1 / math.sqrt(fan_in)
+            init.uniform_(self.bias, -bound, bound)
+
+    def forward(self, input: Tensor) -> Tensor:
+        return F.linear(input, self.weight, self.bias)
+
+    def extra_repr(self) -> str:
+        return 'in_features={}, out_features={}, bias={}'.format(
+            self.in_features, self.out_features, self.bias is not None
+        )
+
+class SuperMLP(nn.Module):
   # MLP: FC -> Activation -> Drop -> FC -> Drop
   def __init__(self, in_features, hidden_features: Optional[int] = None,
                out_features: Optional[int] = None,

lib/spaces/__init__.py

@@ -6,3 +6,5 @@
 
 from .basic_space import Categorical
 from .basic_space import Continuous
+from .basic_op import has_categorical
+from .basic_op import has_continuous

lib/spaces/basic_op.py

@@ -0,0 +1,16 @@
+from spaces.basic_space import Space
+from spaces.basic_space import _EPS
+
+
+def has_categorical(space_or_value, x):
+    if isinstance(space_or_value, Space):
+        return space_or_value.has(x)
+    else:
+        return space_or_value == x
+
+
+def has_continuous(space_or_value, x):
+    if isinstance(space_or_value, Space):
+        return space_or_value.has(x)
+    else:
+        return abs(space_or_value - x) <= _EPS

lib/spaces/basic_space.py

@@ -4,28 +4,65 @@
 
 import abc
 import math
+import copy
 import random
+import numpy as np
 
 from typing import Optional
 
+_EPS = 1e-9
+
 
 class Space(metaclass=abc.ABCMeta):
+    """Basic search space describing the set of possible candidate values for hyperparameter.
+    All search space must inherit from this basic class.
+    """
+
     @abc.abstractmethod
     def random(self, recursion=True):
         raise NotImplementedError
 
+    @abc.abstractproperty
+    def determined(self):
+        raise NotImplementedError
+
     @abc.abstractmethod
     def __repr__(self):
         raise NotImplementedError
 
+    @abc.abstractmethod
+    def has(self, x):
+        """Check whether x is in this search space."""
+        assert not isinstance(
+            x, Space
+        ), "The input value itself can not be a search space."
+
+    def copy(self):
+        return copy.deepcopy(self)
+
 
 class Categorical(Space):
+    """A space contains the categorical values.
+    It can be a nested space, which means that the candidate in this space can also be a search space.
+    """
+
     def __init__(self, *data, default: Optional[int] = None):
         self._candidates = [*data]
         self._default = default
-        assert self._default is None or 0 <= self._default < len(self._candidates), "default >= {:}".format(
-            len(self._candidates)
-        )
+        assert self._default is None or 0 <= self._default < len(
+            self._candidates
+        ), "default >= {:}".format(len(self._candidates))
+        assert len(self) > 0, "Please provide at least one candidate"
+
+    @property
+    def determined(self):
+        if len(self) == 1:
+            return (
+                not isinstance(self._candidates[0], Space)
+                or self._candidates[0].determined
+            )
+        else:
+            return False
 
     def __getitem__(self, index):
         return self._candidates[index]
@@ -38,6 +75,15 @@ class Categorical(Space):
             name=self.__class__.__name__, cs=self._candidates, default=self._default
         )
 
+    def has(self, x):
+        super().has(x)
+        for candidate in self._candidates:
+            if isinstance(candidate, Space) and candidate.has(x):
+                return True
+            elif candidate == x:
+                return True
+        return False
+
     def random(self, recursion=True):
         sample = random.choice(self._candidates)
         if recursion and isinstance(sample, Space):
@@ -46,12 +92,35 @@ class Categorical(Space):
             return sample
 
 
+np_float_types = (np.float16, np.float32, np.float64)
+np_int_types = (
+    np.uint8,
+    np.int8,
+    np.uint16,
+    np.int16,
+    np.uint32,
+    np.int32,
+    np.uint64,
+    np.int64,
+)
+
+
 class Continuous(Space):
-    def __init__(self, lower: float, upper: float, default: Optional[float] = None, log: bool = False):
+    """A space contains the continuous values."""
+
+    def __init__(
+        self,
+        lower: float,
+        upper: float,
+        default: Optional[float] = None,
+        log: bool = False,
+        eps: float = _EPS,
+    ):
         self._lower = lower
         self._upper = upper
         self._default = default
         self._log_scale = log
+        self._eps = eps
 
     @property
     def lower(self):
@@ -65,6 +134,10 @@ class Continuous(Space):
     def default(self):
         return self._default
 
+    @property
+    def determined(self):
+        return abs(self.lower - self.upper) <= self._eps
+
     def __repr__(self):
         return "{name:}(lower={lower:}, upper={upper:}, default_value={default:}, log_scale={log:})".format(
             name=self.__class__.__name__,
@@ -74,6 +147,23 @@ class Continuous(Space):
             log=self._log_scale,
         )
 
+    def convert(self, x):
+        if isinstance(x, np_float_types) and x.size == 1:
+            return float(x), True
+        elif isinstance(x, np_int_types) and x.size == 1:
+            return float(x), True
+        elif isinstance(x, int):
+            return float(x), True
+        elif isinstance(x, float):
+            return float(x), True
+        else:
+            return None, False
+
+    def has(self, x):
+        super().has(x)
+        converted_x, success = self.convert(x)
+        return success and self.lower <= converted_x <= self.upper
+
     def random(self, recursion=True):
         del recursion
         if self._log_scale:

lib/trade_models/naive_v1_model.py

@@ -1,6 +1,8 @@
 ##################################################
 # Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2021 #
 ##################################################
+# Use noise as prediction                        #
+##################################################
 from __future__ import division
 from __future__ import print_function
 
@@ -27,7 +29,11 @@ class NAIVE_V1(Model):
         self.d_feat = d_feat
         self.seed = seed
 
-        self.logger.info("NAIVE-V1 parameters setting: d_feat={:}, seed={:}".format(self.d_feat, self.seed))
+        self.logger.info(
+            "NAIVE-V1 parameters setting: d_feat={:}, seed={:}".format(
+                self.d_feat, self.seed
+            )
+        )
 
         if self.seed is not None:
             random.seed(self.seed)
@@ -49,7 +55,9 @@ class NAIVE_V1(Model):
 
     def model(self, x):
         num = len(x)
-        return np.random.normal(loc=self._mean, scale=self._std, size=num).astype(x.dtype)
+        return np.random.normal(loc=self._mean, scale=self._std, size=num).astype(
+            x.dtype
+        )
 
     def fit(self, dataset: DatasetH):
         def _prepare_dataset(df_data):
@@ -71,9 +79,15 @@ class NAIVE_V1(Model):
         # df_train['feature']['CLOSE1'].values
         # train_dataset['features'][:, -1]
         masks = ~np.isnan(train_dataset["labels"])
-        self._mean, self._std = np.mean(train_dataset["labels"][masks]), np.std(train_dataset["labels"][masks])
-        train_mse_loss = self.mse(self.model(train_dataset["features"]), train_dataset["labels"])
-        valid_mse_loss = self.mse(self.model(valid_dataset["features"]), valid_dataset["labels"])
+        self._mean, self._std = np.mean(train_dataset["labels"][masks]), np.std(
+            train_dataset["labels"][masks]
+        )
+        train_mse_loss = self.mse(
+            self.model(train_dataset["features"]), train_dataset["labels"]
+        )
+        valid_mse_loss = self.mse(
+            self.model(valid_dataset["features"]), valid_dataset["labels"]
+        )
         self.logger.info("Training MSE loss: {:}".format(train_mse_loss))
         self.logger.info("Validation MSE loss: {:}".format(valid_mse_loss))
         self.fitted = True

lib/trade_models/naive_v2_model.py

@@ -29,7 +29,11 @@ class NAIVE_V2(Model):
         self.d_feat = d_feat
         self.seed = seed
 
-        self.logger.info("NAIVE parameters setting: d_feat={:}, seed={:}".format(self.d_feat, self.seed))
+        self.logger.info(
+            "NAIVE parameters setting: d_feat={:}, seed={:}".format(
+                self.d_feat, self.seed
+            )
+        )
 
         if self.seed is not None:
             random.seed(self.seed)
@@ -79,8 +83,12 @@ class NAIVE_V2(Model):
         )
         # df_train['feature']['CLOSE1'].values
         # train_dataset['features'][:, -1]
-        train_mse_loss = self.mse(self.model(train_dataset["features"]), train_dataset["labels"])
-        valid_mse_loss = self.mse(self.model(valid_dataset["features"]), valid_dataset["labels"])
+        train_mse_loss = self.mse(
+            self.model(train_dataset["features"]), train_dataset["labels"]
+        )
+        valid_mse_loss = self.mse(
+            self.model(valid_dataset["features"]), valid_dataset["labels"]
+        )
         self.logger.info("Training MSE loss: {:}".format(train_mse_loss))
         self.logger.info("Validation MSE loss: {:}".format(valid_mse_loss))
         self.fitted = True

lib/trade_models/quant_transformer.py

@@ -37,14 +37,22 @@ from qlib.data.dataset.handler import DataHandlerLP
 
 
 DEFAULT_OPT_CONFIG = dict(
-    epochs=200, lr=0.001, batch_size=2000, early_stop=20, loss="mse", optimizer="adam", num_workers=4
+    epochs=200,
+    lr=0.001,
+    batch_size=2000,
+    early_stop=20,
+    loss="mse",
+    optimizer="adam",
+    num_workers=4,
 )
 
 
 class QuantTransformer(Model):
     """Transformer-based Quant Model"""
 
-    def __init__(self, net_config=None, opt_config=None, metric="", GPU=0, seed=None, **kwargs):
+    def __init__(
+        self, net_config=None, opt_config=None, metric="", GPU=0, seed=None, **kwargs
+    ):
         # Set logger.
         self.logger = get_module_logger("QuantTransformer")
         self.logger.info("QuantTransformer PyTorch version...")
@@ -53,7 +61,9 @@ class QuantTransformer(Model):
         self.net_config = net_config or DEFAULT_NET_CONFIG
         self.opt_config = opt_config or DEFAULT_OPT_CONFIG
         self.metric = metric
-        self.device = torch.device("cuda:{:}".format(GPU) if torch.cuda.is_available() and GPU >= 0 else "cpu")
+        self.device = torch.device(
+            "cuda:{:}".format(GPU) if torch.cuda.is_available() and GPU >= 0 else "cpu"
+        )
         self.seed = seed
 
         self.logger.info(
@@ -84,11 +94,17 @@ class QuantTransformer(Model):
         self.logger.info("model size: {:.3f} MB".format(count_parameters(self.model)))
 
         if self.opt_config["optimizer"] == "adam":
-            self.train_optimizer = optim.Adam(self.model.parameters(), lr=self.opt_config["lr"])
+            self.train_optimizer = optim.Adam(
+                self.model.parameters(), lr=self.opt_config["lr"]
+            )
         elif self.opt_config["optimizer"] == "adam":
-            self.train_optimizer = optim.SGD(self.model.parameters(), lr=self.opt_config["lr"])
+            self.train_optimizer = optim.SGD(
+                self.model.parameters(), lr=self.opt_config["lr"]
+            )
         else:
-            raise NotImplementedError("optimizer {:} is not supported!".format(optimizer))
+            raise NotImplementedError(
+                "optimizer {:} is not supported!".format(optimizer)
+            )
 
         self.fitted = False
         self.model.to(self.device)
@@ -111,7 +127,9 @@ class QuantTransformer(Model):
         else:
             raise ValueError("unknown metric `{:}`".format(self.metric))
 
-    def train_or_test_epoch(self, xloader, model, loss_fn, metric_fn, is_train, optimizer=None):
+    def train_or_test_epoch(
+        self, xloader, model, loss_fn, metric_fn, is_train, optimizer=None
+    ):
         if is_train:
             model.train()
         else:
@@ -173,7 +191,11 @@ class QuantTransformer(Model):
         )
 
         save_dir = get_or_create_path(save_dir, return_dir=True)
-        self.logger.info("Fit procedure for [{:}] with save path={:}".format(self.__class__.__name__, save_dir))
+        self.logger.info(
+            "Fit procedure for [{:}] with save path={:}".format(
+                self.__class__.__name__, save_dir
+            )
+        )
 
         def _internal_test(ckp_epoch=None, results_dict=None):
             with torch.no_grad():
@@ -186,8 +208,10 @@ class QuantTransformer(Model):
                 test_loss, test_score = self.train_or_test_epoch(
                     test_loader, self.model, self.loss_fn, self.metric_fn, False, None
                 )
-                xstr = "train-score={:.6f}, valid-score={:.6f}, test-score={:.6f}".format(
-                    train_score, valid_score, test_score
+                xstr = (
+                    "train-score={:.6f}, valid-score={:.6f}, test-score={:.6f}".format(
+                        train_score, valid_score, test_score
+                    )
                 )
                 if ckp_epoch is not None and isinstance(results_dict, dict):
                     results_dict["train"][ckp_epoch] = train_score
@@ -199,18 +223,26 @@ class QuantTransformer(Model):
         ckp_path = os.path.join(save_dir, "{:}.pth".format(self.__class__.__name__))
         if os.path.exists(ckp_path):
             ckp_data = torch.load(ckp_path)
-            stop_steps, best_score, best_epoch = ckp_data['stop_steps'], ckp_data['best_score'], ckp_data['best_epoch']
-            start_epoch, best_param = ckp_data['start_epoch'], ckp_data['best_param']
-            results_dict = ckp_data['results_dict']
-            self.model.load_state_dict(ckp_data['net_state_dict'])
-            self.train_optimizer.load_state_dict(ckp_data['opt_state_dict'])
+            stop_steps, best_score, best_epoch = (
+                ckp_data["stop_steps"],
+                ckp_data["best_score"],
+                ckp_data["best_epoch"],
+            )
+            start_epoch, best_param = ckp_data["start_epoch"], ckp_data["best_param"]
+            results_dict = ckp_data["results_dict"]
+            self.model.load_state_dict(ckp_data["net_state_dict"])
+            self.train_optimizer.load_state_dict(ckp_data["opt_state_dict"])
             self.logger.info("Resume from existing checkpoint: {:}".format(ckp_path))
         else:
             stop_steps, best_score, best_epoch = 0, -np.inf, -1
             start_epoch, best_param = 0, None
-            results_dict = dict(train=OrderedDict(), valid=OrderedDict(), test=OrderedDict())
+            results_dict = dict(
+                train=OrderedDict(), valid=OrderedDict(), test=OrderedDict()
+            )
             _, eval_str = _internal_test(-1, results_dict)
-            self.logger.info("Training from scratch, metrics@start: {:}".format(eval_str))
+            self.logger.info(
+                "Training from scratch, metrics@start: {:}".format(eval_str)
+            )
 
         for iepoch in range(start_epoch, self.opt_config["epochs"]):
             self.logger.info(
@@ -219,20 +251,35 @@ class QuantTransformer(Model):
                 )
             )
             train_loss, train_score = self.train_or_test_epoch(
-                train_loader, self.model, self.loss_fn, self.metric_fn, True, self.train_optimizer
+                train_loader,
+                self.model,
+                self.loss_fn,
+                self.metric_fn,
+                True,
+                self.train_optimizer,
+            )
+            self.logger.info(
+                "Training :: loss={:.6f}, score={:.6f}".format(train_loss, train_score)
             )
-            self.logger.info("Training :: loss={:.6f}, score={:.6f}".format(train_loss, train_score))
 
             current_eval_scores, eval_str = _internal_test(iepoch, results_dict)
             self.logger.info("Evaluating :: {:}".format(eval_str))
 
             if current_eval_scores["valid"] > best_score:
-                stop_steps, best_epoch, best_score = 0, iepoch, current_eval_scores["valid"]
+                stop_steps, best_epoch, best_score = (
+                    0,
+                    iepoch,
+                    current_eval_scores["valid"],
+                )
                 best_param = copy.deepcopy(self.model.state_dict())
             else:
                 stop_steps += 1
                 if stop_steps >= self.opt_config["early_stop"]:
-                    self.logger.info("early stop at {:}-th epoch, where the best is @{:}".format(iepoch, best_epoch))
+                    self.logger.info(
+                        "early stop at {:}-th epoch, where the best is @{:}".format(
+                            iepoch, best_epoch
+                        )
+                    )
                     break
             save_info = dict(
                 net_config=self.net_config,
@@ -247,9 +294,11 @@ class QuantTransformer(Model):
                 start_epoch=iepoch + 1,
             )
             torch.save(save_info, ckp_path)
-        self.logger.info("The best score: {:.6f} @ {:02d}-th epoch".format(best_score, best_epoch))
+        self.logger.info(
+            "The best score: {:.6f} @ {:02d}-th epoch".format(best_score, best_epoch)
+        )
         self.model.load_state_dict(best_param)
-        _, eval_str = _internal_test('final', results_dict)
+        _, eval_str = _internal_test("final", results_dict)
         self.logger.info("Reload the best parameter :: {:}".format(eval_str))
 
         if self.use_gpu:

lib/trade_models/transformers.py

@@ -33,7 +33,15 @@ DEFAULT_NET_CONFIG = dict(
 
 
 class Attention(nn.Module):
-    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0.0, proj_drop=0.0):
+    def __init__(
+        self,
+        dim,
+        num_heads=8,
+        qkv_bias=False,
+        qk_scale=None,
+        attn_drop=0.0,
+        proj_drop=0.0,
+    ):
         super(Attention, self).__init__()
         self.num_heads = num_heads
         head_dim = dim // num_heads
@@ -46,8 +54,16 @@ class Attention(nn.Module):
 
     def forward(self, x):
         B, N, C = x.shape
-        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
-        q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)
+        qkv = (
+            self.qkv(x)
+            .reshape(B, N, 3, self.num_heads, C // self.num_heads)
+            .permute(2, 0, 3, 1, 4)
+        )
+        q, k, v = (
+            qkv[0],
+            qkv[1],
+            qkv[2],
+        )  # make torchscript happy (cannot use tensor as tuple)
 
         attn = (q @ k.transpose(-2, -1)) * self.scale
         attn = attn.softmax(dim=-1)
@@ -76,13 +92,25 @@ class Block(nn.Module):
         super(Block, self).__init__()
         self.norm1 = norm_layer(dim)
         self.attn = Attention(
-            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=mlp_drop
+            dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            qk_scale=qk_scale,
+            attn_drop=attn_drop,
+            proj_drop=mlp_drop,
         )
         # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
-        self.drop_path = xlayers.DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+        self.drop_path = (
+            xlayers.DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+        )
         self.norm2 = norm_layer(dim)
         mlp_hidden_dim = int(dim * mlp_ratio)
-        self.mlp = xlayers.MLP(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=mlp_drop)
+        self.mlp = xlayers.MLP(
+            in_features=dim,
+            hidden_features=mlp_hidden_dim,
+            act_layer=act_layer,
+            drop=mlp_drop,
+        )
 
     def forward(self, x):
         x = x + self.drop_path(self.attn(self.norm1(x)))
@@ -144,9 +172,13 @@ class TransformerModel(nn.Module):
         self.input_embed = SimpleEmbed(d_feat, embed_dim=embed_dim)
 
         self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
-        self.pos_embed = xlayers.PositionalEncoder(d_model=embed_dim, max_seq_len=max_seq_len, dropout=pos_drop)
+        self.pos_embed = xlayers.PositionalEncoder(
+            d_model=embed_dim, max_seq_len=max_seq_len, dropout=pos_drop
+        )
 
-        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+        dpr = [
+            x.item() for x in torch.linspace(0, drop_path_rate, depth)
+        ]  # stochastic depth decay rule
         self.blocks = nn.ModuleList(
             [
                 Block(
@@ -184,7 +216,9 @@ class TransformerModel(nn.Module):
         batch, flatten_size = x.shape
         feats = self.input_embed(x)  # batch * 60 * 64
 
-        cls_tokens = self.cls_token.expand(batch, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
+        cls_tokens = self.cls_token.expand(
+            batch, -1, -1
+        )  # stole cls_tokens impl from Phil Wang, thanks
         feats_w_ct = torch.cat((cls_tokens, feats), dim=1)
         feats_w_tp = self.pos_embed(feats_w_ct)
 

tests/test_basic.py

@@ -1,38 +0,0 @@
-#####################################################
-# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2021.03 #
-#####################################################
-import sys
-import unittest
-import pytest
-from pathlib import Path
-
-lib_dir = (Path(__file__).parent / ".." / "lib").resolve()
-print("library path: {:}".format(lib_dir))
-if str(lib_dir) not in sys.path:
-    sys.path.insert(0, str(lib_dir))
-
-from spaces import Categorical
-from spaces import Continuous
-
-
-class TestBasicSpace(unittest.TestCase):
-    def test_categorical(self):
-        space = Categorical(1, 2, 3, 4)
-        for i in range(4):
-            self.assertEqual(space[i], i + 1)
-        self.assertEqual("Categorical(candidates=[1, 2, 3, 4], default_index=None)", str(space))
-
-    def test_continuous(self):
-        space = Continuous(0, 1)
-        self.assertGreaterEqual(space.random(), 0)
-        self.assertGreaterEqual(1, space.random())
-
-        lower, upper = 1.5, 4.6
-        space = Continuous(lower, upper, log=False)
-        values = []
-        for i in range(100000):
-            x = space.random()
-            self.assertGreaterEqual(x, lower)
-            self.assertGreaterEqual(upper, x)
-            values.append(x)
-        self.assertAlmostEqual((lower + upper) / 2, sum(values) / len(values), places=2)

tests/test_basic_space.py

@@ -0,0 +1,81 @@
+#####################################################
+# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2021.03 #
+#####################################################
+import sys, random
+import unittest
+import pytest
+from pathlib import Path
+
+lib_dir = (Path(__file__).parent / ".." / "lib").resolve()
+print("library path: {:}".format(lib_dir))
+if str(lib_dir) not in sys.path:
+    sys.path.insert(0, str(lib_dir))
+
+from spaces import Categorical
+from spaces import Continuous
+
+
+class TestBasicSpace(unittest.TestCase):
+    """Test the basic search spaces."""
+
+    def test_categorical(self):
+        space = Categorical(1, 2, 3, 4)
+        for i in range(4):
+            self.assertEqual(space[i], i + 1)
+        self.assertEqual(
+            "Categorical(candidates=[1, 2, 3, 4], default_index=None)", str(space)
+        )
+
+    def test_continuous(self):
+        random.seed(999)
+        space = Continuous(0, 1)
+        self.assertGreaterEqual(space.random(), 0)
+        self.assertGreaterEqual(1, space.random())
+
+        lower, upper = 1.5, 4.6
+        space = Continuous(lower, upper, log=False)
+        values = []
+        for i in range(1000000):
+            x = space.random()
+            self.assertGreaterEqual(x, lower)
+            self.assertGreaterEqual(upper, x)
+            values.append(x)
+        self.assertAlmostEqual((lower + upper) / 2, sum(values) / len(values), places=2)
+        self.assertEqual(
+            "Continuous(lower=1.5, upper=4.6, default_value=None, log_scale=False)",
+            str(space),
+        )
+
+    def test_determined_and_has(self):
+        # Test Non-nested Space
+        space = Categorical(1, 2, 3, 4)
+        self.assertFalse(space.determined)
+        self.assertTrue(space.has(2))
+        self.assertFalse(space.has(6))
+        space = Categorical(4)
+        self.assertTrue(space.determined)
+
+        space = Continuous(0.11, 0.12)
+        self.assertTrue(space.has(0.115))
+        self.assertFalse(space.has(0.1))
+        self.assertFalse(space.determined)
+        space = Continuous(0.11, 0.11)
+        self.assertTrue(space.determined)
+
+        # Test Nested Space
+        space_1 = Categorical(1, 2, 3, 4)
+        space_2 = Categorical(1)
+        nested_space = Categorical(space_1)
+        self.assertFalse(nested_space.determined)
+        self.assertTrue(nested_space.has(4))
+        nested_space = Categorical(space_2)
+        self.assertTrue(nested_space.determined)
+
+        # Test Nested Space 2
+        nested_space = Categorical(
+            Categorical(1, 2, 3),
+            Categorical(4, Categorical(5, 6, 7, Categorical(8, 9), 10), 11),
+            12,
+        )
+        for i in range(1, 13):
+            self.assertTrue(nested_space.has(i))