moospread 0.1.0__py3-none-any.whl

moospread/tasks/re_torch.py

@@ -0,0 +1,394 @@
+import torch
+import torch.nn as nn
+from pymoo.problems import get_problem
+from moospread.problem import PymooProblemTorch
+import os
+import numpy as np
+
+######## RE Problems ########
+# Note: For the sake of differentiability, we will use the strict bounds:
+# - Lower bound: xl + 1e-6 (instead of xl)
+# - Upper bound: xu - 1e-6 (instead of xu)
+# ref_point: The default reference points are suitable for the "online" mode.
+
+class RE21(PymooProblemTorch):
+    def __init__(self, path = None, ref_point=None, **kwargs):
+        F = 10.0
+        sigma = 10.0
+        tmp_val = F / sigma  # = 1.
+        super().__init__(n_var=4, n_obj=2, 
+                         xl=torch.tensor([tmp_val, 
+                                          (2.0**0.5)*tmp_val, 
+                                          (2.0**0.5)*tmp_val, 
+                                          tmp_val]) + 1e-6,
+                         xu=torch.full((4,), 
+                                       3.0*tmp_val) - 1e-6,
+                         vtype=float, **kwargs)
+        if ref_point is None:
+            self.ref_point = [3144.44, 0.05]
+        else:
+            self.ref_point = ref_point
+            
+        self.path = path
+    
+    def _calc_pareto_front(self, n_pareto_points: int = 100) -> torch.Tensor:
+        assert self.path is not None, "Path to Pareto front file not specified."
+        front = np.loadtxt(self.path)
+        return torch.from_numpy(front).to(self.device)
+
+    def _evaluate(self, X: torch.Tensor, out: dict, *args, **kwargs) -> None:
+        """
+        X: (N, 4) with columns [x1, x2, x3, x4]
+        Returns F: (N, 2) with [f0, f1]
+        """
+        # ensure batch dimension
+        if X.dim() == 1:
+            X = X.unsqueeze(0)
+
+        if X.size(-1) != 4:
+            raise ValueError("X must be a (N, 4) tensor.")
+
+        x1, x2, x3, x4 = X.unbind(dim=1)
+
+        # constants
+        F = 10.0
+        sigma = 10.0
+        E = 2.0e5
+        L = 200.0
+        sqrt2 = torch.sqrt(torch.tensor(2.0, dtype=X.dtype, device=X.device))
+
+        # objectives (vectorized)
+        f0 = L * (2.0 * x1 + sqrt2 * x2 + torch.sqrt(x3) + x4)
+
+        # small clamp for numerical safety (bounds already keep > 0)
+        eps = 1e-12
+        x1c = torch.clamp(x1, min=eps)
+        x2c = torch.clamp(x2, min=eps)
+        x3c = torch.clamp(x3, min=eps)
+        x4c = torch.clamp(x4, min=eps)
+
+        f1 = ((F * L) / E) * (
+            (2.0 / x1c) + (2.0 * sqrt2 / x2c) - (2.0 * sqrt2 / x3c) + (2.0 / x4c)
+        )
+
+        out["F"] = torch.stack([f0, f1], dim=1)  # (N, 2)
+        
+
+class RE33(PymooProblemTorch):
+    def __init__(self, path=None, ref_point=None, **kwargs):
+        super().__init__(n_var=4, n_obj=3, 
+                         xl=torch.tensor([55.0, 75.0, 1000.0, 11.0])+ 1e-6,
+                         xu=torch.tensor([80.0, 110.0, 3000.0, 20.0])- 1e-6,
+                         vtype=float, **kwargs)
+        if ref_point is None:
+            self.ref_point = [5.01, 9.84, 4.30]
+        else:
+            self.ref_point = ref_point
+            
+        self.path = path
+    
+    def _calc_pareto_front(self, n_pareto_points: int = 100) -> torch.Tensor:
+        assert self.path is not None, "Path to Pareto front file not specified."
+        front = np.loadtxt(self.path)
+        return torch.from_numpy(front).to(self.device)
+
+
+    def _evaluate(self, X: torch.Tensor, out: dict, *args, **kwargs) -> None:
+        """
+        X: (N, 4) with columns [x1, x2, x3, x4]
+        Returns: (N, 3) with [f0, f1, f2], where f2 is sum of constraint violations.
+        """
+        if X.dim() != 2 or X.size(-1) != 4:
+            raise ValueError("X must be a (N, 4) tensor.")
+
+        x1, x2, x3, x4 = X.unbind(dim=1)
+
+        # ----- Objectives -----
+        # f0 = 4.9e-5 * (x2^2 - x1^2) * (x4 - 1)
+        d2 = x2**2 - x1**2
+        f0 = 4.9e-5 * d2 * (x4 - 1.0)
+
+        # f1 = (9.82e6) * (x2^2 - x1^2) / (x3 * x4 * (x2^3 - x1^3))
+        d3 = x2**3 - x1**3
+        eps = torch.tensor(1e-12, dtype=X.dtype, device=X.device)
+
+        def safe_signed(x):
+            s = torch.where(x >= 0, 1.0, -1.0)
+            return s * torch.clamp(x.abs(), min=eps)
+
+        denom_f1 = (torch.clamp(x3, min=eps) *
+                    torch.clamp(x4, min=eps) *
+                    safe_signed(d3))
+        f1 = (9.82e6) * d2 / denom_f1
+
+        # ----- Constraints (g >= 0 means satisfied) -----
+        # g0 = (x2 - x1) - 20
+        g0 = (x2 - x1) - 20.0
+
+        # g1 = 0.4 - x3 / (pi * (x2^2 - x1^2))
+        denom_g1 = torch.pi * safe_signed(d2)
+        g1 = 0.4 - (x3 / denom_g1)
+
+        # g2 = 1 - (2.22e-3 * x3 * (x2^3 - x1^3)) / (x2^2 - x1^2)^2
+        denom_g2 = safe_signed(d2)**2
+        g2 = 1.0 - (2.22e-3 * x3 * d3) / denom_g2
+
+        # g3 = (2.66e-2 * x3 * x4 * (x2^3 - x1^3)) / (x2^2 - x1^2) - 900
+        denom_g3 = safe_signed(d2)
+        g3 = (2.66e-2 * x3 * x4 * d3) / denom_g3 - 900.0
+
+        Gs = torch.stack([g0, g1, g2, g3], dim=1)  # (N, 4)
+        # violation = -g if g < 0 else 0
+        violations = torch.where(Gs < 0, -Gs, torch.zeros_like(Gs))
+        f2 = violations.sum(dim=1)
+
+        out["F"] = torch.stack([f0, f1, f2], dim=1)  # (N, 3)
+
+
+
+class RE34(PymooProblemTorch):
+    def __init__(self, path = None, ref_point=None, **kwargs):
+        super().__init__(n_var=5, n_obj=3, 
+                         xl=torch.full((5,), 1.0) + 1e-6,
+                         xu=torch.full((5,), 3.0) - 1e-6,
+                         vtype=float, **kwargs)
+        if ref_point is None:
+            self.ref_point = [1.86472022e+03, 1.18199394e+01, 2.90399938e-01]
+        else:
+            self.ref_point = ref_point
+            
+        self.path = path
+
+    def _calc_pareto_front(self, n_pareto_points: int = 100) -> torch.Tensor:
+        assert self.path is not None, "Path to Pareto front file not specified."
+        front = np.loadtxt(self.path)
+        return torch.from_numpy(front).to(self.device)
+
+    def _evaluate(self, x: torch.Tensor, out: dict, *args, **kwargs) -> None:
+        # ensure batch dimension
+        if x.dim() == 1:
+            x = x.unsqueeze(0)                      # [5] → [1,5]
+
+        # unpack variables
+        x1 = x[:, 0]
+        x2 = x[:, 1]
+        x3 = x[:, 2]
+        x4 = x[:, 3]
+        x5 = x[:, 4]
+
+        # prepare output
+        f = torch.zeros((x.shape[0], self.n_obj),
+                        dtype=x.dtype, device=x.device)
+
+        # objective 1
+        f[:, 0] = (
+              1640.2823
+            + 2.3573285 * x1
+            + 2.3220035 * x2
+            + 4.5688768 * x3
+            + 7.7213633 * x4
+            + 4.4559504 * x5
+        )
+
+        # objective 2
+        f[:, 1] = (
+              6.5856
+            + 1.15    * x1
+            - 1.0427  * x2
+            + 0.9738  * x3
+            + 0.8364  * x4
+            - 0.3695  * x1 * x4
+            + 0.0861  * x1 * x5
+            + 0.3628  * x2 * x4
+            - 0.1106  * x1 * x1
+            - 0.3437  * x3 * x3
+            + 0.1764  * x4 * x4
+        )
+
+        # objective 3
+        f[:, 2] = (
+             -0.0551
+            + 0.0181 * x1
+            + 0.1024 * x2
+            + 0.0421 * x3
+            - 0.0073 * x1 * x2
+            + 0.0240 * x2 * x3
+            - 0.0118 * x2 * x4
+            - 0.0204 * x3 * x4
+            - 0.0080 * x3 * x5
+            - 0.0241 * x2 * x2
+            + 0.0109 * x4 * x4
+        )
+
+        # if input was a single vector, return shape (3,)
+        out["F"] = f
+
+       
+class RE37(PymooProblemTorch):
+    "Rocket Injector Design (RE37)."
+
+    def __init__(self, path=None, ref_point=None, **kwargs):
+        super().__init__(n_var=4, n_obj=3, 
+                         xl=torch.zeros(4) + 1e-6,
+                         xu=torch.ones(4) - 1e-6,
+                         vtype=float, **kwargs)
+        if ref_point is None:
+            self.ref_point = [1.1022, 1.20726899, 1.20318656]
+        else:
+            self.ref_point = ref_point
+            
+        self.path = path
+
+    def _calc_pareto_front(self, n_pareto_points: int = 100) -> torch.Tensor:
+        assert self.path is not None, "Path to Pareto front file not specified."
+        front = np.loadtxt(self.path)
+        return torch.from_numpy(front).to(self.device)
+
+    def _evaluate(self, x: torch.Tensor, out: dict, *args, **kwargs) -> None:
+        # Unpack features: each is (B,)
+        xAlpha, xHA, xOA, xOPTT = x.unbind(dim=1)
+
+        # Compute objectives vectorized over batch
+        f1 = (
+            0.692
+            + 0.477 * xAlpha
+            - 0.687 * xHA
+            - 0.080 * xOA
+            - 0.0650 * xOPTT
+            - 0.167  * xAlpha**2
+            - 0.0129 * xHA * xAlpha
+            + 0.0796 * xHA**2
+            - 0.0634 * xOA * xAlpha
+            - 0.0257 * xOA * xHA
+            + 0.0877 * xOA**2
+            - 0.0521 * xOPTT * xAlpha
+            + 0.00156 * xOPTT * xHA
+            + 0.00198 * xOPTT * xOA
+            + 0.0184  * xOPTT**2
+        )
+
+        f2 = (
+            0.153
+            - 0.322  * xAlpha
+            + 0.396  * xHA
+            + 0.424  * xOA
+            + 0.0226 * xOPTT
+            + 0.175  * xAlpha**2
+            + 0.0185 * xHA * xAlpha
+            - 0.0701 * xHA**2
+            - 0.251  * xOA * xAlpha
+            + 0.179  * xOA * xHA
+            + 0.0150 * xOA**2
+            + 0.0134 * xOPTT * xAlpha
+            + 0.0296 * xOPTT * xHA
+            + 0.0752 * xOPTT * xOA
+            + 0.0192 * xOPTT**2
+        )
+
+        f3 = (
+            0.370
+            - 0.205  * xAlpha
+            + 0.0307 * xHA
+            + 0.108  * xOA
+            + 1.019  * xOPTT
+            - 0.135  * xAlpha**2
+            + 0.0141 * xHA * xAlpha
+            + 0.0998 * xHA**2
+            + 0.208  * xOA * xAlpha
+            - 0.0301 * xOA * xHA
+            - 0.226  * xOA**2
+            + 0.353  * xOPTT * xAlpha
+            - 0.0497 * xOPTT * xOA
+            - 0.423  * xOPTT**2
+            + 0.202  * xHA * xAlpha**2
+            - 0.281  * xOA * xAlpha**2
+            - 0.342  * xHA**2 * xAlpha
+            - 0.245  * xHA**2 * xOA
+            + 0.281  * xOA**2 * xHA
+            - 0.184  * xOPTT**2 * xAlpha
+            - 0.281  * xHA * xAlpha * xOA
+        )
+
+        # Stack into (B, 3)
+        out["F"] = torch.stack([f1, f2, f3], dim=1)
+
+
+class RE41(PymooProblemTorch):
+    r"""Car side impact problem (RE41).
+
+    See [Tanabe2020]_ for details.
+    
+    """
+    def __init__(self, path=None, ref_point=None, **kwargs):
+        super().__init__(n_var=7, n_obj=4, 
+                         xl=torch.tensor([0.5, 0.45, 0.5, 0.5, 0.875, 0.4, 0.4], dtype=torch.float) + 1e-6,
+                         xu=torch.tensor([1.5, 1.35, 1.5, 1.5, 2.625, 1.2, 1.2], dtype=torch.float) - 1e-6,
+                         vtype=float, **kwargs)
+        if ref_point is None:
+            self.ref_point = [47.04480682,  4.86997366, 14.40049127, 10.3941957 ]
+        else:
+            self.ref_point = ref_point
+            
+        self.path = path
+
+        self.continuous_inds = list(range(7))
+
+    def _evaluate(self, X: torch.Tensor, out: dict, *args, **kwargs) -> None:
+        X1, X2, X3, X4, X5, X6, X7 = torch.split(X, 1, -1)
+        f1 = (
+            1.98
+            + 4.9 * X1
+            + 6.67 * X2
+            + 6.98 * X3
+            + 4.01 * X4
+            + 1.78 * X5
+            + 10**-5 * X6
+            + 2.73 * X7
+        )
+        f2 = 4.72 - 0.5 * X4 - 0.19 * X2 * X3
+        V_MBP = 10.58 - 0.674 * X1 * X2 - 0.67275 * X2
+        V_FD = 16.45 - 0.489 * X3 * X7 - 0.843 * X5 * X6
+        f3 = 0.5 * (V_MBP + V_FD)
+        g1 = 1 - 1.16 + 0.3717 * X2 * X4 + 0.0092928 * X3
+        g2 = (
+            0.32
+            - 0.261
+            + 0.0159 * X1 * X2
+            + 0.06486 * X1
+            + 0.019 * X2 * X7
+            - 0.0144 * X3 * X5
+            - 0.0154464 * X6
+        )
+        g3 = (
+            0.32
+            - 0.214
+            - 0.00817 * X5
+            + 0.045195 * X1
+            + 0.0135168 * X1
+            - 0.03099 * X2 * X6
+            + 0.018 * X2 * X7
+            - 0.007176 * X3
+            - 0.023232 * X3
+            + 0.00364 * X5 * X6
+            + 0.018 * X2.pow(2)
+        )
+        g4 = 0.32 - 0.74 + 0.61 * X2 + 0.031296 * X3 + 0.031872 * X7 - 0.227 * X2.pow(2)
+        g5 = 32 - 28.98 - 3.818 * X3 + 4.2 * X1 * X2 - 1.27296 * X6 + 2.68065 * X7
+        g6 = (
+            32
+            - 33.86
+            - 2.95 * X3
+            + 5.057 * X1 * X2
+            + 3.795 * X2
+            + 3.4431 * X7
+            - 1.45728
+        )
+        g7 = 32 - 46.36 + 9.9 * X2 + 4.4505 * X1
+        g8 = 4 - f2
+        g9 = 9.9 - V_MBP
+        g10 = 15.7 - V_FD
+        g = torch.cat([g1, g2, g3, g4, g5, g6, g7, g8, g9, g10], dim=-1)
+        zero = torch.tensor(0.0, dtype=X.dtype, device=X.device)
+        g = torch.where(g < 0, -g, zero)
+        f4 = g.sum(dim=-1, keepdim=True)
+        out["F"] = torch.cat([f1, f2, f3, f4], dim=-1)

moospread/tasks/zdt_torch.py

@@ -0,0 +1,112 @@
+import torch
+import torch.nn as nn
+from pymoo.problems import get_problem
+from moospread.problem import PymooProblemTorch
+import numpy as np
+
+######## ZDT 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 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)

moospread/utils/__init__.py

@@ -0,0 +1,8 @@
+from moospread.utils.ditmoo import DiTMOO
+from moospread.utils.lhs import LHS
+from moospread.utils.spread_utils import get_ddpm_dataloader, is_pass_function
+from moospread.utils.mobo_utils import *
+from moospread.utils.offline_utils import *
+from moospread.utils.misc import * 
+from moospread.utils.constraint_utils import PMGDASolver
+

moospread/utils/constraint_utils/__init__.py

@@ -0,0 +1,2 @@
+from moospread.utils.constraint_utils.prefs import get_uniform_pref
+from moospread.utils.constraint_utils.pmgda_solver import PMGDASolver

moospread/utils/constraint_utils/gradient.py

@@ -0,0 +1,72 @@
+import torch
+
+
+def get_moo_Jacobian_batch(x_batch, y_batch, n_obj):
+    '''
+        Input : x_batch: (batch_size, n_var)
+                y_batch: (batch_size, n_obj)
+                n_obj: int
+        Return: grad_batch: (batch_size, n_obj, n_var)
+    '''
+    grad_batch = []
+    batch_size = len(x_batch)
+    for batch_idx in range(batch_size):
+        # grad_arr.append(get_moo_Jacobian(xs, ys, n_obj))
+        grad_arr_idx = [0,] * n_obj
+        for obj_idx in range(n_obj):
+            y_batch[batch_idx][obj_idx].backward(retain_graph = True)
+            grad_arr_idx[obj_idx] = x_batch.grad[batch_idx].clone()
+            x_batch.grad.zero_()
+        grad_batch.append( torch.stack(grad_arr_idx) )
+    return torch.stack(grad_batch)
+
+def get_moo_Jacobian(x, y, n_obj):
+    grad_arr = [0] * n_obj
+    for obj_idx in range(n_obj):
+        y[obj_idx].backward(retain_graph=True)
+        grad_arr[obj_idx] = x.grad.clone()
+        x.grad.zero_()
+    grad_arr = torch.stack(grad_arr)
+    return grad_arr
+
+def flatten_grads(grads_dict):
+    return torch.cat( [v.view(-1) for _, v in grads_dict.items()] )
+
+def calc_gradients_mtl(data, batch, model, objectives):
+    # store gradients and objective values
+    gradients = []
+    obj_values = []
+    for i, objective in enumerate(objectives):
+        model.zero_grad()
+        logits = model(data)
+        output = objective(logits['logits'], **batch)
+        output.backward()
+        obj_values.append(output.item())
+        gradients.append({})
+        private_params = model.private_params() if hasattr(model, 'private_params') else []
+        for name, param in model.named_parameters():
+            not_private = all([p not in name for p in private_params])
+            if not_private and param.requires_grad and param.grad is not None:
+                gradients[i][name] = param.grad.data.detach().clone()
+    return gradients
+
+def get_grads_from_model(loss, model):
+    G = [0,] * len(loss)
+    for idx, l in enumerate(loss):
+        model.zero_grad()
+        l.backward(retain_graph=True)
+        G[idx] = get_flatten_grad(model)
+    return torch.stack(G)
+
+def get_flatten_grad(model):
+    grad = []
+    for param in model.parameters():
+        if param.grad is not None:
+            grad.append(param.grad.view(-1))
+        else:
+            grad.append(torch.zeros_like(param.view(-1)))
+    grad = torch.cat(grad)
+    return grad
+
+def numel_params(model):
+    return sum(p.numel() for p in model.parameters() if p.requires_grad)

moospread/utils/constraint_utils/mgda_core.py

@@ -0,0 +1,69 @@
+import torch
+from numpy import array
+import numpy as np
+from cvxopt import matrix, solvers
+solvers.options['show_progress'] = False
+# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+
+def solve_mgda_analy(grad_1, grad_2):
+    '''
+        Solve_mgda_analy only support 2-objective case.
+        grad_i.shape: (n,).
+        This function support grad_i as both Tensor and numpy.
+    '''
+    v1v1 = grad_1 @ grad_1
+    v2v2 = grad_2 @ grad_2
+    v1v2 = grad_1 @ grad_2
+
+    if v1v2 >= v1v1:
+        gamma = 0.999
+    elif v1v2 >= v2v2:
+        gamma = 0.001
+    else:
+        gamma = -1.0 * ((v1v2 - v2v2) / (v1v1 + v2v2 - 2 * v1v2))
+    coeff = torch.Tensor([gamma, 1 - gamma] )
+    # gw = coeff[0] * grad_1 + coeff[1] * grad_2
+    # else:
+    #     return gw
+    return coeff
+
+
+
+
+
+def solve_mgda(Jacobian):
+    '''
+        Input Jacobian: (m,n).
+        Output alpha: (m,)
+    '''
+
+    # m : n_obj
+    # n : n_var
+
+    m = Jacobian.shape[0]
+    if m == 2:
+        return solve_mgda_analy(Jacobian[0], Jacobian[1])
+    else:
+        Q = (Jacobian @ Jacobian.T).cpu().detach().numpy()
+
+        Q = matrix(np.float64(Q))
+        p = np.zeros(m)
+        A = np.ones(m)
+
+        A = matrix(A, (1, m))
+        b = matrix(1.0)
+
+        G_cvx = -np.eye(m)
+        h = [0.0] * m
+        h = matrix(h)
+
+        G_cvx = matrix(G_cvx)
+        p = matrix(p)
+        sol = solvers.qp(Q, p, G_cvx, h, A, b)
+
+        res = np.array(sol['x']).squeeze()
+        alpha = res / sum(res)  # important. Does res already satisfy sum=1?
+        return alpha
+
+