update NSGA-test

DTLZ_problem/__init__.py

@@ -1 +1 @@
-from .dataset import create_dataset
+from .dataset import create_dataset, evaluate, get_pf, get_moea_data

DTLZ_problem/dataset.py

@@ -1,8 +1,14 @@
+from __future__ import annotations
+
 from typing import Tuple, List
 
 import numpy as np
 from pymoo.core.problem import Problem
 from pymoo.problems.many.dtlz import get_ref_dirs
+from pymoo.util.ref_dirs import get_reference_directions
+from pymoo.optimize import minimize
+from pymoo.indicators.igd import IGD
+from pymoo.algorithms.moo.nsga3 import NSGA3
 
 
 class DTLZb(Problem):
@@ -44,12 +50,15 @@ def obj_func(self, X_, g, alpha=1):
 
 class DTLZ1b(DTLZb):
     def __init__(self, n_var=7, n_obj=3, delta1=0, delta2=0, **kwargs):
+        self.delta1 = delta1
+        self.delta2 = delta2
         super().__init__(n_var, n_obj, delta1, delta2, **kwargs)
 
     def _calc_pareto_front(self, ref_dirs=None):
         if ref_dirs is None:
             ref_dirs = get_ref_dirs(self.n_obj)
-        return 0.5 * ref_dirs
+        coefficient = max(0.5, 0.5 * (100 + self.delta1) * self.delta2)
+        return coefficient * ref_dirs
 
     def obj_func(self, X_, g):
         f = []
@@ -69,6 +78,19 @@ def _evaluate(self, x, out, *args, **kwargs):
         out["F"] = self.obj_func(X_, g)
 
 
+def pf_data(n_var: int, n_objective: int, delta1: int, delta2: int) -> np.ndarray:
+    problem = DTLZ1b(n_var=n_var, n_obj=n_objective, delta1=delta1, delta2=delta2)  # change delta here
+    ref_dirs = get_reference_directions("das-dennis", n_objective, n_partitions=12)
+    N = ref_dirs.shape[0]
+    # create the algorithm object
+    algorithm = NSGA3(pop_size=N, ref_dirs=ref_dirs)
+    # execute the optimization
+    res = minimize(problem,
+                   algorithm,
+                   termination=('n_gen', 100))
+    return res.X
+
+
 def create_dataset_inner(x, n_dim: Tuple[int, int], delta: Tuple[List[int], List[int]]) -> Tuple[
     np.ndarray, np.ndarray
 ]:
@@ -99,10 +121,12 @@ def create_dataset_inner(x, n_dim: Tuple[int, int], delta: Tuple[List[int], List
     return new_x, y
 
 
-def create_dataset(problem_dim: Tuple[int, int], x=None, n_problem=None, spt_qry=None, delta=None, normalize_targets=True, **_) -> Tuple[
-    Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray],
-    Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]
-]:
+def create_dataset(problem_dim: Tuple[int, int], x=None, n_problem=None, spt_qry=None, delta=None,
+                   normalize_targets=True, **_) -> Tuple[
+    Tuple[
+        Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray],
+        Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]
+    ], Tuple[float | None, float | None]]:
     """
     Parameters
     ----------
@@ -111,8 +135,8 @@ def create_dataset(problem_dim: Tuple[int, int], x=None, n_problem=None, spt_qry
     x : Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]
         [x_spt_train, x_qry_train, x_spt_test, x_qry_test]
         The input data, shape (4, n_problem, n_spt, n_variables)
-    delta : Tuple[train_spt_delta, train_qry_delta, test_spt_delta, test_qry_delta]
-        The delta values, shape (4, 2, n_problem)
+    delta : Tuple[train_delta,  test_delta]
+        The delta values, shape (2, 2, n_problem)
     n_problem : Tuple[int, int]
         [n_train, n_test]
     spt_qry : Tuple[int, int]
@@ -123,35 +147,48 @@ def create_dataset(problem_dim: Tuple[int, int], x=None, n_problem=None, spt_qry
     Returns
     -------
     Tuple[
-        Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray],
-        Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]
+        Tuple[
+            Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray],
+            Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]
+        ],
+        Tuple[float | None, float | None]
     ]
-        The first element is the training set [support set, support label, query set, query label]
-        The second element is the test set [support set, support label, query set, query label]
+        In the first Tuple:
+            The first element is the training set [support set, support label, query set, query label]
+            The second element is the test set [support set, support label, query set, query label]
+        In the second Tuple:
+            The minimum and maximum of the targets (to be used for normalization)
     """
     n_var, n_obj = problem_dim
-    if x is not None:
-        assert len(x) == 4
-    else:
-        # generate x
-        x = []
-        for i in range(2):
-            for j in range(2):
-                x.append(np.random.rand(n_problem[i], spt_qry[j], n_var))
 
     if delta is not None:
-        assert len(x) == 4
+        assert len(delta) == 2
     else:
         # generate delta
         delta = []
+        for i in range(2):
+            delta1 = np.random.randint(0, 100, n_problem[i])
+            delta2 = np.random.randint(0, 10, n_problem[i])
+            delta.append([delta1, delta2])
+
+    if x is not None:
+        assert len(x) == 4
+    else:
+        # generate x
+        x = []
         for i in range(2):
             for j in range(2):
-                delta1 = np.random.randint(0, 100, n_problem[i])
-                delta2 = np.random.randint(0, 10, n_problem[i])
-                delta.append([delta1, delta2])
+                x_pf = []
+                for k in range(n_problem[i]):
+                    ps = pf_data(n_var, n_obj, delta[i][j][k], delta[i][j][k])
+                    x_pf.append(ps[np.random.choice(ps.shape[0], int(0.5 * spt_qry[j]))])
+                x_pf = np.array(x_pf)
+                x_ran = np.random.rand(n_problem[i], spt_qry[j] - int(0.5 * spt_qry[j]), n_var)
+                x.append(np.concatenate((x_pf, x_ran), axis=1))
+                # x.append(np.random.rand(n_problem[i], spt_qry[j], n_var))
 
-    train_set = [*create_dataset_inner(x[0], problem_dim, delta[0]), *create_dataset_inner(x[1], problem_dim, delta[1])]
-    test_set = [*create_dataset_inner(x[2], problem_dim, delta[2]), *create_dataset_inner(x[3], problem_dim, delta[3])]
+    train_set = [*create_dataset_inner(x[0], problem_dim, delta[0]), *create_dataset_inner(x[1], problem_dim, delta[0])]
+    test_set = [*create_dataset_inner(x[2], problem_dim, delta[1]), *create_dataset_inner(x[3], problem_dim, delta[1])]
     if normalize_targets:
         minimum = np.min(np.concatenate([train_set[1], test_set[1]]), axis=None)
         train_set[1] -= minimum
@@ -163,13 +200,123 @@ def create_dataset(problem_dim: Tuple[int, int], x=None, n_problem=None, spt_qry
         test_set[1] /= maximum
         train_set[3] /= maximum
         test_set[3] /= maximum
-    return tuple(train_set), tuple(test_set)
+    else:
+        maximum = None
+        minimum = None
+    return (tuple(train_set), tuple(test_set)), (minimum, maximum)
+
+
+def evaluate(x: np.ndarray, delta: Tuple[int, int],
+             n_objectives: int, min_max: Tuple[float | None, float | None]) -> np.ndarray:
+    """
+    Parameters
+    ----------
+    x : np.ndarray
+        The input data, shape (n_point, n_variables)
+    delta : Tuple[int, int]
+        The delta1 and delta2
+    n_objectives: int
+        The number of objectives
+    min_max : Tuple[float | None, float | None]
+        The minimum and maximum of the targets, if None, the targets will not be normalized
+
+    Returns
+    -------
+    np.ndarray
+        The output data, shape (n_point, n_objectives)
+    """
+    n_variables = x.shape[1]
+    problem = DTLZ1b(n_var=n_variables, n_obj=n_objectives, delta1=delta[0], delta2=delta[1])
+    y = problem.evaluate(x)
+    if min_max[0] is not None:
+        y -= min_max[0]
+        y /= min_max[1]
+    return y
+
+
+def get_pf(n_var: int, n_objectives: int, delta: Tuple[int, int],
+           min_max: Tuple[float | None, float | None]) -> np.ndarray:
+    """
+    Parameters
+    ----------
+    n_var: int
+        The number of variables
+    n_objectives: int
+        The number of objectives
+    delta : Tuple[int, int]
+        The delta1 and delta2
+    min_max : Tuple[float | None, float | None]
+        The minimum and maximum of the targets, if None, the targets will not be normalized
+
+    Returns
+    -------
+    np.ndarray
+        The parato front, shape (n_point, n_objectives)
+    """
+    problem = DTLZ1b(n_var=n_var, n_obj=n_objectives, delta1=delta[0], delta2=delta[1])
+    ref_dirs = get_reference_directions("das-dennis", n_objectives, n_partitions=12)
+    pf = problem.pareto_front(ref_dirs)
+    if min_max[0] is not None:
+        pf -= min_max[0]
+        pf /= min_max[1]
+    return pf
+
+
+def get_moea_data(n_var: int, n_objectives: int, delta: Tuple[int, int], algorithm, n_gen: int,
+                  min_max: Tuple[float | None, float | None]) -> Tuple[
+    np.ndarray, list, list
+]:
+    """
+    Parameters
+    ----------
+    n_var: int
+        The number of variables
+    n_objectives: int
+        The number of objectives
+    delta: Tuple[int, int]
+        The delta1 and delta2
+    algorithm: 
+        MOEA algorithm
+    n_gen: int
+        number of generation
+    min_max: Tuple[float | None, float | None]
+        The minimum and maximum of the targets, if None, the targets will not be normalized
+
+    Returns
+    -------
+    Tuple[np.ndarray, list, list]
+        moea_pf: The parato front, shape (n_point, n_objectives)
+        n_evals: The number of function evaluations
+        igd: The IGD
+    """
+    problem = DTLZ1b(n_var=n_var, n_obj=n_objectives, delta1=delta[0], delta2=delta[1])  # change delta here
+    res = minimize(problem,
+                   algorithm,
+                   termination=('n_gen', n_gen),
+                   save_history=True,
+                   verbose=False)
+    moea_pf = res.F
+
+    hist = res.history
+    hist_F, n_evals = [], []
+    for algo in hist:
+        n_evals.append(algo.evaluator.n_eval)
+        opt = algo.opt
+        feas = np.where(opt.get("feasible"))[0]
+        hist_F.append(opt.get("F")[feas])
+    metric = IGD(moea_pf, zero_to_one=True)
+    igd = [metric.do(_F) for _F in hist_F]
+
+    if min_max[0] is not None:
+        moea_pf -= min_max[0]
+        moea_pf /= min_max[1]
+    return moea_pf, n_evals, igd
 
 
 def test():
-    n_dim = (7, 3)
+    n_dim = (8, 3)
     train_set, test_set = create_dataset(n_dim, n_problem=(4, 2), spt_qry=(5, 20))  # (12, 5, 7)
-    print(train_set[0].shape, train_set[1].shape, train_set[2].shape, train_set[3].shape)
+    print(train_set[0].shape, train_set[1].shape, test_set[0].shape, test_set[1].shape)
 
 
 if __name__ == '__main__':

examples/__init__.py

@@ -1,2 +1,3 @@
 from examples import example
 from examples import example_sinewave
+from .baseline_nn import train as train_baseline_nn

examples/baseline_nn.py

@@ -0,0 +1,77 @@
+from typing import Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+class BaselineNn(nn.Module):
+    def __init__(self, n_dim: Tuple[int, int]):
+        super(BaselineNn, self).__init__()
+        self.n_args_in, self.n_args_out = n_dim
+        self.layers = nn.Sequential(
+            nn.Linear(self.n_args_in, 2 * self.n_args_in),
+            nn.ReLU(),
+            nn.Linear(2 * self.n_args_in, 4 * self.n_args_in),
+            nn.ReLU(),
+            nn.Linear(4 * self.n_args_in, 4 * self.n_args_in),
+            nn.ReLU(),
+            nn.Linear(4 * self.n_args_in, 2 * self.n_args_in),
+            nn.ReLU(),
+            nn.Linear(2 * self.n_args_in, self.n_args_out),
+        )
+
+    def forward(self, x, return_tensor=False):
+        if isinstance(x, np.ndarray):
+            x = torch.from_numpy(x).float().to(self.args.device)
+        self.to(self.args.device)
+        if len(x.shape) == 1:
+            x = x.reshape(1, -1)
+        ret = self.layers(x)
+        return ret if return_tensor else ret.detach().cpu().numpy()
+
+
+class Wrapper:
+    def __init__(self, shape):
+        self.shape = shape
+        self.models = [BaselineNn((shape[0], 1)) for _ in range(shape[1])]
+
+    def __call__(self, x, *args, **kwargs):
+        if isinstance(x, np.ndarray):
+            x = torch.from_numpy(x).float().to(self.args.device)
+        for m in self.models:
+            m.__dict__['args'] = self.args
+        ys = []
+        for i in range(self.shape[1]):
+            ys.append(self.models[i](x).flatten()[0])
+        ret = np.array(ys)
+        if 'return_tensor' in kwargs and kwargs['return_tensor']:
+            return torch.from_numpy(ret).float().to(self.args.device)
+        return ret
+
+
+__model = None
+
+
+def train(x: torch.Tensor, y: torch.Tensor, shape: Tuple[int, int],
+          init: bool = False, lr: float = 0.001, n_epochs: int = 1000):
+    global __model
+    if init:
+        __model = Wrapper(shape)
+        return __model
+    for i, s in enumerate(__model.models):
+        optimizer = torch.optim.Adam(s.parameters(), lr=lr)
+        loss_fn = torch.nn.MSELoss()
+        y_tensor = torch.from_numpy(y[i]).float().to(s.args.device)
+        loss_list = []
+        for _ in range(n_epochs):
+            optimizer.zero_grad()
+            y_pred = s(x, return_tensor=True)
+            loss = loss_fn(y_pred, y_tensor)
+            loss_list.append(loss.item())
+            loss.backward()
+            optimizer.step()
+
+        print(f'Loss: {loss_list[-1]}, avg loss: {np.mean(loss_list)}+-{np.std(loss_list)}')
+
+    return __model

examples/example.py

@@ -12,14 +12,14 @@
 def get_args():
     args = NamedDict()
     args.problem_dim = (8, 3)
-    args.train_test = (20, 1)
+    args.train_test = (15, 1)
     args.epoch = 50
     args.update_lr = 0.01
     args.meta_lr = 0.001
-    args.k_spt = 10
-    args.k_qry = 100
-    args.update_step = 5
-    args.update_step_test = 8
+    args.k_spt = 100
+    args.k_qry = 200
+    args.update_step = 30
+    args.update_step_test = 50
     args.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     # args.device = torch.device('cpu')
     return args
@@ -37,6 +37,15 @@ def get_network_structure(args):
         ('linear', [2 * n_args, 4 * n_args]),
         ('relu', [True]),
         ('linear', [1, 2 * n_args]),
+        # ('linear', [100, n_args]),
+        # ('relu', [True]),
+        # ('linear', [200, 100]),
+        # ('relu', [True]),
+        # ('linear', [200, 200]),
+        # ('relu', [True]),
+        # ('linear', [100, 200]),
+        # ('relu', [True]),
+        # ('linear', [1, 100]),
     ]
     return config