moospread 0.1.0__py3-none-any.whl

moospread/problem.py

@@ -0,0 +1,193 @@
+"""
+This module defines an abstract base class for 
+multi-objective optimization problems using PyTorch.
+"""
+
+import numpy as np
+import torch
+from abc import ABC, abstractmethod
+from pymoo.util.cache import Cache
+
+
+def default_shape(problem, n):
+    n_var = problem.n_var
+    return {
+        'F': (n, problem.n_obj),
+        'G': (n, problem.n_ieq_constr),
+        'H': (n, problem.n_eq_constr),
+        'dF': (n, problem.n_obj, n_var),
+        'dG': (n, problem.n_ieq_constr, n_var),
+        'dH': (n, problem.n_eq_constr, n_var),
+        'CV': (n, 1),
+    }
+
+class PymooProblemTorch(ABC):
+    def __init__(
+        self,
+        n_var: int = -1,
+        n_obj: int = 1,
+        n_ieq_constr: int = 0,
+        n_eq_constr: int = 0,
+        xl=None,
+        xu=None,
+        device: torch.device = None,
+        elementwise: bool = False,
+        replace_nan_by: float = None,
+        strict: bool = True,
+        **kwargs
+    ):
+        self.n_var = n_var
+        self.n_obj = n_obj
+        self.n_ieq_constr = n_ieq_constr
+        self.n_eq_constr = n_eq_constr
+        self.device = device or torch.device('cpu')
+        self.elementwise = elementwise
+        self.replace_nan_by = replace_nan_by
+        self.strict = strict
+        self.ref_point = None
+
+        # bounds
+        if n_var > 0 and xl is not None:
+            self.xl = (
+                xl
+                if isinstance(xl, torch.Tensor)
+                else torch.full((n_var,), xl, dtype=torch.float, device=self.device)
+            )
+        else:
+            self.xl = None
+        if n_var > 0 and xu is not None:
+            self.xu = (
+                xu
+                if isinstance(xu, torch.Tensor)
+                else torch.full((n_var,), xu, dtype=torch.float, device=self.device)
+            )
+        else:
+            self.xu = None
+
+        self.is_discrete = False
+        self.is_sequence = False
+        self.original_bounds = None
+        self.need_repair = False
+        if self.has_bounds():
+            self.original_bounds = [self.xl.clone(), self.xu.clone()]
+            self.global_clamping = False
+            self.need_repair = True
+
+    def has_bounds(self) -> bool:
+        return self.xl is not None and self.xu is not None
+
+    def has_constraints(self) -> bool:
+        return (self.n_ieq_constr + self.n_eq_constr) > 0
+
+    def name(self) -> str:
+        return self.__class__.__name__
+    
+    def bounds(self):
+        return [self.xl, self.xu]
+    
+    def evaluate(
+        self,
+        x: torch.Tensor,
+        return_as_dict: bool = False,
+        return_values_of=None
+    ):
+        """
+        Evaluate objectives on input tensor x.
+
+        Args:
+            x: Tensor of shape (batch_size, n_var) with requires_grad=True.
+            return_as_dict: If True, return the full dict of computed values.
+            return_values_of: Optional list of keys from the internal out dict to return alongside F.
+
+        Returns:
+            If return_as_dict=True: the out dict (including all computed values).
+            Else if return_values_of provided: a tuple (F, dict_of_requested_values).
+            Else: Tensor F of shape (batch_size, n_obj).
+        """
+        if not isinstance(x, torch.Tensor):
+            raise TypeError("Input must be a torch.Tensor")
+        x.data = torch.nan_to_num(x.data, nan=0.0, posinf=1.0, neginf=0.0)
+
+        # Ensure is is within bounds
+        if self.has_bounds():
+            if self.xl is not None:
+                x.data = torch.max(x.data, self.xl.to(x.device))
+            if self.xu is not None:
+                x.data = torch.min(x.data, self.xu.to(x.device))
+
+        out: dict = {}
+        self._evaluate(x, out)
+
+        # Prepare base return
+        F = out.get("F")
+        # If specific extra values requested
+        extra = None
+        if return_values_of is not None:
+            if isinstance(return_values_of, (list, tuple)):
+                extra = {k: out[k] for k in return_values_of if k in out}
+            else:
+                key = return_values_of
+                extra = {key: out[key]} if key in out else {}
+
+        if return_as_dict:
+            if extra is not None:
+                out.update(extra)
+            return out
+
+        if extra is not None:
+            if len(extra) == 1:
+                return list(extra.values())[0]
+            return tuple(extra[k] for k in extra)
+
+        return F
+
+    @abstractmethod
+    def _evaluate(self, x, out: dict):
+        """
+        User-implemented evaluation. 
+        If elementwise: x is 1d Tensor of shape (n_var,).
+        Otherwise x is 2d Tensor of shape (n_samples, n_var).
+        out should be populated with keys 'F', 'G', 'H' as torch.Tensors.
+        """
+        pass
+    
+    @Cache
+    def pareto_front(self, *args, **kwargs):
+        pf = self._calc_pareto_front(*args, **kwargs)
+        if pf is None:
+            return None
+        if not isinstance(pf, torch.Tensor):
+            pf = torch.as_tensor(pf, dtype=torch.float32)
+        if pf.dim() == 1:
+            pf = pf.unsqueeze(0)
+        if pf.size(1) == 2:
+            pf = pf[pf[:, 0].argsort()]
+        return pf
+
+    def __str__(self):
+        return (
+            f"Problem(name={self.name()}, n_var={self.n_var}, n_obj={self.n_obj}, "
+            f"n_ieq_constr={self.n_ieq_constr}, n_eq_constr={self.n_eq_constr})"
+        )
+
+
+class BaseProblem(PymooProblemTorch):
+    """
+    Base Problem class to ensure PyTorch differentiability.
+    """
+
+    def __init__(self, n_var=30, n_obj=2, **kwargs):
+        super().__init__(n_var=n_var, n_obj=n_obj, 
+                         xl=torch.zeros(n_var, dtype=torch.float) + 1e-6,
+                         xu=torch.ones(n_var, dtype=torch.float) - 1e-6,
+                         vtype=float, **kwargs)
+        
+    def _calc_pareto_front(self, 
+                           n_pareto_points: int = 100) -> torch.Tensor:
+        pass
+
+    def _evaluate(self, x: torch.Tensor, 
+                  out: dict, 
+                  *args, 
+                  **kwargs) -> None:
+        pass

moospread/tasks/__init__.py

@@ -0,0 +1,4 @@
+from moospread.tasks.dtlz_torch import DTLZ, DTLZ2, DTLZ4, DTLZ7
+from moospread.tasks.zdt_torch import ZDT, ZDT1, ZDT2, ZDT3
+from moospread.tasks.re_torch import RE21, RE33, RE34, RE37, RE41
+from moospread.tasks.mw_torch import MW7

moospread/tasks/dtlz_torch.py

@@ -0,0 +1,139 @@
+import torch
+from pymoo.problems import get_problem
+from moospread.problem import PymooProblemTorch
+import math
+from torch import Tensor
+
+######## DTLZ Problems ########
+# Note: For the sake of differentiability, we will use the strict bounds:
+# - Lower bound: 0 + 1e-6 (instead of 0)
+# - Upper bound: 1 - 1e-6 (instead of 1)
+# ref_point: The default reference points are suitable for the "online" mode.
+
+class DTLZ(PymooProblemTorch):
+    r"""Base class for DTLZ problems.
+
+    See [Deb2005dtlz]_ for more details on DTLZ.
+    """
+
+    def __init__(self, n_var=30, n_obj=3, **kwargs):
+        super().__init__(n_var=n_var, n_obj=n_obj,
+                         xl=torch.zeros(n_var, dtype=torch.float) + 1e-6,
+                         xu=torch.ones(n_var, dtype=torch.float) - 1e-6,
+                         vtype=float, **kwargs)
+
+        if n_var <= n_obj:
+            raise ValueError(
+                f"n_var must be > n_obj, but got {n_var} and {n_obj}."
+            )
+        self.continuous_inds = list(range(n_var))
+        self.k = n_var - n_obj + 1
+       
+        
+class DTLZ2(DTLZ):
+    r"""DLTZ2 test problem.
+
+    d-dimensional problem evaluated on `[0, 1]^d`:
+
+        f_0(x) = (1 + g(x)) * cos(x_0 * pi / 2)
+        f_1(x) = (1 + g(x)) * sin(x_0 * pi / 2)
+        g(x) = \sum_{i=m}^{d-1} (x_i - 0.5)^2
+
+    The pareto front is given by the unit hypersphere \sum{i} f_i^2 = 1.
+    Note: the pareto front is completely concave. The goal is to minimize
+    both objectives.
+    """
+    def __init__(self, n_var=30, n_obj=3, ref_point=None, **kwargs):
+        super().__init__(n_var, n_obj, **kwargs)
+        if ref_point is None and n_obj == 3:
+            self.ref_point = [2.8390, 2.9011, 2.8575]
+        else:
+            assert ref_point is not None, "Please provide a reference point for n_obj != 3"
+            self.ref_point = ref_point
+    
+    def pareto_front(self):
+        return get_problem("dtlz2", n_var=self.n_var, n_obj=self.n_obj).pareto_front()
+
+    def _evaluate(self, X: torch.Tensor, out: dict, *args, **kwargs) -> None:
+        X_m = X[..., -self.k :]
+        g_X = (X_m - 0.5).pow(2).sum(dim=-1)
+        g_X_plus1 = 1 + g_X
+        fs = []
+        pi_over_2 = math.pi / 2
+        for i in range(self.n_obj):
+            idx = self.n_obj - 1 - i
+            f_i = g_X_plus1.clone()
+            f_i *= torch.cos(X[..., :idx] * pi_over_2).prod(dim=-1)
+            if i > 0:
+                f_i *= torch.sin(X[..., idx] * pi_over_2)
+            fs.append(f_i)
+        out["F"] = torch.stack(fs, dim=-1)
+
+        
+class DTLZ4(DTLZ):
+    
+    def __init__(self, n_var=30, n_obj=3, ref_point=None, **kwargs):
+        super().__init__(n_var, n_obj, **kwargs)
+        if ref_point is None and n_obj == 3:
+            self.ref_point = [3.2675, 2.6443, 2.4263]
+        else:
+            assert ref_point is not None, "Please provide a reference point for n_obj != 3"
+            self.ref_point = ref_point
+    
+    def pareto_front(self):
+        return get_problem("dtlz4", n_var=self.n_var, n_obj=self.n_obj).pareto_front()
+
+    def _evaluate(self, x: torch.Tensor, out: dict, *args, **kwargs) -> None:
+        X_, X_M = x[:, :self.n_obj - 1], x[:, self.n_obj - 1:]
+        alpha  = 100
+        g_X = (X_M - 0.5).pow(2).sum(dim=-1)
+        g_X_plus1 = 1 + g_X
+        fs = []
+        pi_over_2 = math.pi / 2
+        for i in range(self.n_obj):
+            idx = self.n_obj - 1 - i
+            f_i = g_X_plus1.clone()
+            f_i *= torch.cos((X_[..., :idx]**alpha) * pi_over_2).prod(dim=-1)
+            if i > 0:
+                f_i *= torch.sin((X_[..., idx]**alpha) * pi_over_2)
+            fs.append(f_i)
+        out["F"] = torch.stack(fs, dim=-1)
+        
+ 
+class DTLZ7(DTLZ):
+    r"""DTLZ7 test problem.
+
+    d-dimensional problem evaluated on `[0, 1]^d`:
+        f_0(x) = x_0
+        f_1(x) = x_1
+        ...
+        f_{M-1}(x) = (1 + g(X_m)) * h(f_0, f_1, ..., f_{M-2}, g, x)
+        h(f_0, f_1, ..., f_{M-2}, g, x) =
+        M - sum_{i=0}^{M-2} f_i(x)/(1+g(x)) * (1 + sin(3 * pi * f_i(x)))
+
+    This test problem has 2M-1 disconnected Pareto-optimal regions in the search space.
+
+    The pareto frontier corresponds to X_m = 0.
+    """
+    def __init__(self, n_var=30, n_obj=3, ref_point=None, **kwargs):
+        super().__init__(n_var, n_obj, **kwargs)
+        if ref_point is None and n_obj == 3:
+            self.ref_point = [0.9984, 0.9961, 22.8114]
+        else:
+            assert ref_point is not None, "Please provide a reference point for n_obj != 3"
+            self.ref_point = ref_point
+    
+    def pareto_front(self):
+        return get_problem("dtlz7", n_var=self.n_var, n_obj=self.n_obj).pareto_front()
+
+    def _evaluate(self, X: torch.Tensor, out: dict, *args, **kwargs) -> None:
+        f = []
+        for i in range(0, self.n_obj - 1):
+            f.append(X[..., i])
+        f = torch.stack(f, dim=-1)
+
+        g_X = 1 + 9 / self.k * torch.sum(X[..., -self.k :], dim=-1)
+        h = self.n_obj - torch.sum(
+            f / (1 + g_X.unsqueeze(-1)) * (1 + torch.sin(3 * math.pi * f)), dim=-1
+        )
+        out["F"] = torch.cat([f, ((1 + g_X) * h).unsqueeze(-1)], dim=-1)

moospread/tasks/mw_torch.py

@@ -0,0 +1,274 @@
+import torch
+import torch.nn as nn
+from pymoo.problems import get_problem
+from pymoo.util.remote import Remote
+from moospread.problem import PymooProblemTorch
+import numpy as np
+
+######## MW Problems ########
+# Note: For the sake of differentiability, we will use the strict bounds:
+# - Lower bound: 0 + 1e-6 (instead of 0)
+# - Upper bound: 1 - 1e-6 (instead of 1)
+# ref_point: The default reference points are suitable for the "online" mode.
+
+import torch
+from torch import Tensor
+from typing import Optional
+
+class MW(PymooProblemTorch):
+    def __init__(self, n_var, n_obj,
+                n_ieq_constr: int = 0,
+                n_eq_constr: int = 0,
+                **kwargs):
+        super().__init__(n_var=n_var, n_obj=n_obj, 
+                         xl=torch.zeros(n_var, dtype=torch.float) + 1e-6,
+                         xu=torch.ones(n_var, dtype=torch.float) - 1e-6,
+                         n_ieq_constr=n_ieq_constr, n_eq_constr=n_eq_constr,
+                         vtype=float, **kwargs)
+        
+    # ---- smooth building blocks (torch versions of LA1/2/3) ----
+    @staticmethod
+    def LA1(A, B, C, D, theta: Tensor) -> Tensor:
+        # A * (sin(B * pi * theta^C))^D
+        t = theta.pow(C)
+        return torch.as_tensor(A, dtype=theta.dtype, device=theta.device) * \
+               torch.sin(torch.as_tensor(B, dtype=theta.dtype, device=theta.device) * torch.pi * t).pow(D)
+
+    @staticmethod
+    def LA2(A, B, C, D, theta: Tensor) -> Tensor:
+        # A * (sin(B * theta^C))^D
+        t = theta.pow(C)
+        return torch.as_tensor(A, dtype=theta.dtype, device=theta.device) * \
+               torch.sin(torch.as_tensor(B, dtype=theta.dtype, device=theta.device) * t).pow(D)
+
+    @staticmethod
+    def LA3(A, B, C, D, theta: Tensor) -> Tensor:
+        # A * (cos(B * theta^C))^D
+        t = theta.pow(C)
+        return torch.as_tensor(A, dtype=theta.dtype, device=theta.device) * \
+               torch.cos(torch.as_tensor(B, dtype=theta.dtype, device=theta.device) * t).pow(D)
+
+    # ---- distance/landscape functions ----
+    def g1(self, X: Tensor) -> Tensor:
+        """
+        X: (batch, n_var) tensor
+        returns: (batch,) tensor
+        """
+        if not isinstance(X, Tensor):
+            X = torch.as_tensor(X, dtype=torch.get_default_dtype())
+        d = self.n_var
+        n = d - self.n_obj
+
+        # z shape: (batch, d - (n_obj-1))
+        z = X[:, self.n_obj - 1:].pow(n)
+
+        # i shape: (d - (n_obj-1),)
+        i = torch.arange(self.n_obj - 1, d, device=X.device, dtype=X.dtype)
+        offset = 0.5 + i / (2.0 * d)  # broadcast over batch
+        delta = z - offset
+        exp_term = 1.0 - torch.exp(-10.0 * (delta * delta))
+        distance = 1.0 + exp_term.sum(dim=1)
+        return distance
+
+    def g2(self, X: Tensor) -> Tensor:
+        """
+        X: (batch, n_var) tensor
+        returns: (batch,) tensor
+        """
+        if not isinstance(X, Tensor):
+            X = torch.as_tensor(X, dtype=torch.get_default_dtype())
+        d = self.n_var
+        n = float(d)
+
+        i = torch.arange(self.n_obj - 1, d, device=X.device, dtype=X.dtype)  # (len,)
+        z = 1.0 - torch.exp(-10.0 * (X[:, self.n_obj - 1:] - i / n).pow(2))
+        contrib = (0.1 / n) * (z * z) + 1.5 - 1.5 * torch.cos(2.0 * torch.pi * z)
+        distance = 1.0 + contrib.sum(dim=1)
+        return distance
+
+    def g3(self, X: Tensor) -> Tensor:
+        """
+        X: (batch, n_var) tensor
+        returns: (batch,) tensor
+        """
+        if not isinstance(X, Tensor):
+            X = torch.as_tensor(X, dtype=torch.get_default_dtype())
+
+        a = X[:, self.n_obj - 1:]                         # last block
+        b = X[:, self.n_obj - 2:-1] - 0.5                 # previous block (shifted)
+        inner = a + (b * b) - 1.0
+        contrib = 2.0 * inner.pow(2)
+        distance = 1.0 + contrib.sum(dim=1)
+        return distance
+
+
+class MW7(MW):
+    def __init__(self, n_var: int = 15, ref_point=None, **kwargs):
+        super().__init__(n_var=n_var, 
+                         n_obj=2, 
+                         n_ieq_constr=2, **kwargs)
+        if ref_point is None:
+            self.ref_point = [2.0, 2.0]
+        else:
+            self.ref_point = ref_point
+        
+    def _calc_pareto_front(self, n_pareto_points: int = 100) -> torch.Tensor:
+        return Remote.get_instance().load("pymoo", "pf", "MW", "MW7.pf")
+
+    @torch.no_grad()
+    def _check_shapes(self, X: Tensor):
+        if X.ndim != 2 or X.shape[1] != self.n_var:
+            raise ValueError(f"X must have shape (batch, {self.n_var})")
+
+    def _evaluate(self, X: torch.Tensor, out: dict, *args, **kwargs) -> None:
+        """
+        X: (batch, n_var) tensor
+        returns: dict with
+            F: (batch, 2)
+            G: (batch, 2)
+        """
+        if not isinstance(X, Tensor):
+            X = torch.as_tensor(X, dtype=torch.get_default_dtype())
+        self._check_shapes(X)
+
+        # g from MW.g3 (>= 1, so safe to divide by)
+        g = self.g3(X)                       # (batch,)
+
+        f0 = g * X[:, 0]                     # (batch,)
+        # 1 - (f0/g)^2, clamped for numerical safety
+        one_minus_sq = 1.0 - (f0 / g).pow(2)
+        f1 = g * torch.sqrt(torch.clamp(one_minus_sq, min=0.0))
+
+        # atan2 handles all quadrants and f0=0 safely (better than arctan(f1/f0))
+        atan = torch.atan2(f1, f0)          # (batch,)
+
+        # Constraints via LA2 (uses sin(B * theta^C)^D)
+        term0 = 1.2 + torch.abs(self.LA2(0.4, 4.0, 1.0, 16.0, atan))
+        g0 = f0.pow(2) + f1.pow(2) - term0.pow(2)
+
+        term1 = 1.15 - self.LA2(0.2, 4.0, 1.0, 8.0, atan)
+        g1 = term1.pow(2) - f0.pow(2) - f1.pow(2)
+
+        F = torch.stack([f0, f1], dim=1)    # (batch, 2)
+        G = torch.stack([g0, g1], dim=1)    # (batch, 2)
+        out["F"] = F
+        out["G"] = G
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+class ZDT(PymooProblemTorch):
+    """
+    Base class to ensure PyTorch differentiability.
+    Provides a default `evaluate` that preserves gradients and guards against NaNs.
+    """
+
+    def __init__(self, n_var=30, **kwargs):
+        super().__init__(n_var=n_var, n_obj=2, 
+                         xl=torch.zeros(n_var, dtype=torch.float) + 1e-6,
+                         xu=torch.ones(n_var, dtype=torch.float) - 1e-6,
+                         vtype=float, **kwargs)
+        
+class ZDT1(ZDT):
+    """
+    ZDT1 test problem in PyTorch, fully differentiable.
+    """
+    def __init__(self, n_var=30, ref_point=None, **kwargs):
+        super().__init__(n_var, **kwargs)
+        if ref_point is None:
+            self.ref_point = [0.9994, 6.0576]
+        else:
+            self.ref_point = ref_point
+    
+    def _calc_pareto_front(self, n_pareto_points: int = 100) -> torch.Tensor:
+        x = torch.linspace(0.0, 1.0, n_pareto_points, device=self.device)
+        return torch.stack([x, 1.0 - torch.sqrt(x)], dim=1)
+
+    def _evaluate(self, x: torch.Tensor, out: dict, *args, **kwargs) -> None:
+        # x: (batch_size, n_var)
+        # Objective f1
+        f1 = x[:, 0]
+        # Auxiliary g
+        g = 1.0 + 9.0 / (self.n_var - 1) * torch.sum(x[:, 1:], dim=1)
+        # Avoid negative division
+        term = torch.clamp(f1 / g, min=0.0)
+        # Objective f2
+        f2 = g * (1.0 - torch.sqrt(term))
+        out["F"] = torch.stack([f1, f2], dim=1)
+
+class ZDT2(ZDT):
+    """
+    ZDT2 test problem in PyTorch, fully differentiable.
+    """
+    def __init__(self, n_var=30, ref_point=None, **kwargs):
+        super().__init__(n_var, **kwargs)
+        if ref_point is None:
+            self.ref_point = [0.9994, 6.8960]
+        else:
+            self.ref_point = ref_point
+            
+    def _calc_pareto_front(self, n_pareto_points: int = 100) -> torch.Tensor:
+        x = torch.linspace(0.0, 1.0, n_pareto_points, device=self.device)
+        return torch.stack([x, 1.0 - x.pow(2)], dim=1)
+
+    def _evaluate(self, x: torch.Tensor, out: dict, *args, **kwargs) -> None:
+        f1 = x[:, 0]
+        c = torch.sum(x[:, 1:], dim=1)
+        g = 1.0 + 9.0 * c / (self.n_var - 1)
+        term = torch.clamp(f1 / g, min=0.0)
+        f2 = g * (1.0 - term.pow(2))
+        out["F"] = torch.stack([f1, f2], dim=1)
+
+class ZDT3(ZDT):
+    """
+    ZDT3 test problem in PyTorch, fully differentiable.
+    """
+    def __init__(self, n_var=30, ref_point=None, **kwargs):
+        super().__init__(n_var, **kwargs)
+        if ref_point is None:
+            self.ref_point = [0.9994, 6.0571]
+        else:
+            self.ref_point = ref_point
+     
+    def _calc_pareto_front(
+        self,
+        n_points: int = 100,
+        flatten: bool = True
+    ) -> torch.Tensor:
+        regions = [
+            [0.0, 0.0830015349],
+            [0.182228780, 0.2577623634],
+            [0.4093136748, 0.4538821041],
+            [0.6183967944, 0.6525117038],
+            [0.8233317983, 0.8518328654]
+        ]
+        pf_list = []
+        points_per_region = int(n_points / len(regions))
+        for r in regions:
+            x1 = torch.linspace(r[0], r[1], points_per_region, device=self.device)
+            x2 = 1.0 - torch.sqrt(x1) - x1 * torch.sin(10.0 * torch.pi * x1)
+            pf_list.append(torch.stack([x1, x2], dim=1))
+        return torch.cat(pf_list, dim=0) if flatten else torch.stack(pf_list, dim=0)
+
+    def _evaluate(self, x: torch.Tensor, out: dict, *args, **kwargs) -> None:
+        f1 = x[:, 0]
+        c = torch.sum(x[:, 1:], dim=1)
+        g = 1.0 + 9.0 * c / (self.n_var - 1)
+        term = torch.clamp(f1 / g, min=0.0)
+        f2 = g * (1.0 - torch.sqrt(term) - term * torch.sin(10.0 * torch.pi * f1))
+        out["F"] = torch.stack([f1, f2], dim=1)