Repository: EMI-Group/evox

Issue:
NaN may occur in the CEC2022 SO benchmark
### Desciption
In the CEC2022 single objective benchmark, F10 ~ F12 all use the `compose_operat` function in `OperatFunc` class, where the codes below may lead to `NaN` items in `term1`, resulting in final `NaN` in the evaluated fitness.
```python
diff_square_sum = jnp.sum((x - Os_mat) ** 2, axis=1)
term1 = 1 / jnp.sqrt(diff_square_sum)
term2 = jnp.exp(-0.5 * diff_square_sum / (sigma**2 * D))
W = term1 * term2
W_norm = W / jnp.sum(W)
```

### Possible Solution
```python
diff_square_sum = jnp.sum((x - Os_mat) ** 2, axis=1)
term1 = 1 / jnp.sqrt(diff_square_sum)
term2 = jnp.exp(-0.5 * diff_square_sum / (sigma**2 * D))
W = term1 * term2
term1_nan = jnp.isnan(term1)
W_norm = jnp.select(jnp.any(term1_nan), term1_nan / jnp.count_nonzero(term1_nan), W / jnp.sum(W))
```

### Location
```
- problems
  - numerical
    - cec2022_so.py
      - line 121
```

Is this the intended behavior of CEC2022 single objective benchmark or a bug?


### Note
The overall code of CEC2022 benchmark is hard to read due to the lack of documentations and improper abbreviations, please fix these problems if possible.

Hints from issue comments:
@MinyangChen: Yes, just as you mentioned, NaN values can indeed be produced unexpectedly, which is something we do not wish to see. Your Possible Solution addresses this issue very well.

CEC2022 does have some areas where readability is poor. In order to keep the numerical results produced by this code consistent with those of the original code, I had to port over some complicated parts. I apologize for that.@BillHuang2001: Please take a look at https://github.com/EMI-Group/evox/pull/131. If no there is objection, I will merge the PR and close this issue.

Solution patch:
diff --git a/docs/source/api/algorithms/mo/im_moea.rst b/docs/source/api/algorithms/mo/im_moea.rst
new file mode 100644
index 00000000..fd4dce24
--- /dev/null
+++ b/docs/source/api/algorithms/mo/im_moea.rst
@@ -0,0 +1,6 @@
+======
+IMMOEA
+======
+
+.. autoclass:: evox.algorithms.IMMOEA
+    :members:
diff --git a/docs/source/api/algorithms/mo/index.rst b/docs/source/api/algorithms/mo/index.rst
index 5f8650f1..4ad3e10c 100644
--- a/docs/source/api/algorithms/mo/index.rst
+++ b/docs/source/api/algorithms/mo/index.rst
@@ -22,3 +22,5 @@ Multi-objective
     spea2
     sra
     tdea
+    rveaa
+    im_moea
diff --git a/docs/source/api/algorithms/mo/rveaa.rst b/docs/source/api/algorithms/mo/rveaa.rst
new file mode 100644
index 00000000..5cb2a656
--- /dev/null
+++ b/docs/source/api/algorithms/mo/rveaa.rst
@@ -0,0 +1,6 @@
+=====
+RVEAa
+=====
+
+.. autoclass:: evox.algorithms.RVEAa
+    :members:
diff --git a/src/evox/algorithms/mo/__init__.py b/src/evox/algorithms/mo/__init__.py
index 4a767383..66e5aa6f 100644
--- a/src/evox/algorithms/mo/__init__.py
+++ b/src/evox/algorithms/mo/__init__.py
@@ -16,3 +16,4 @@ from .tdea import TDEA
 from .bce_ibea import BCEIBEA
 from .lmocso import LMOCSO
 from .im_moea import IMMOEA
+from .rveaa import RVEAa
diff --git a/src/evox/algorithms/mo/eagmoead.py b/src/evox/algorithms/mo/eagmoead.py
index f8e1efd2..fad36be5 100644
--- a/src/evox/algorithms/mo/eagmoead.py
+++ b/src/evox/algorithms/mo/eagmoead.py
@@ -13,6 +13,7 @@
 import jax
 import jax.numpy as jnp
 from functools import partial
+import math
 
 from evox import jit_class, Algorithm, State
 from evox.operators import (
@@ -62,7 +63,7 @@ class EAGMOEAD(Algorithm):
         self.dim = lb.shape[0]
         self.pop_size = pop_size
         self.LGs = LGs
-        self.T = jnp.ceil(self.pop_size / 10).astype(int)
+        self.T = math.ceil(self.pop_size / 10)
 
         self.selection = selection_op
         self.mutation = mutation_op
diff --git a/src/evox/algorithms/mo/rvea.py b/src/evox/algorithms/mo/rvea.py
index 7fdf147d..b41d64a5 100644
--- a/src/evox/algorithms/mo/rvea.py
+++ b/src/evox/algorithms/mo/rvea.py
@@ -92,12 +92,11 @@ class RVEA(Algorithm):
         return state
 
     def ask(self, state):
-        key, subkey, x_key, x1_key, x2_key, mut_key = jax.random.split(state.key, 6)
-        population = state.population
+        key, subkey, x_key, mut_key = jax.random.split(state.key, 4)
 
+        population = state.population
         no_nan_pop = ~jnp.isnan(population).all(axis=1)
         max_idx = jnp.sum(no_nan_pop).astype(int)
-
         pop = population[jnp.where(no_nan_pop, size=self.pop_size, fill_value=-1)]
 
         mating_pool = jax.random.randint(subkey, (self.pop_size,), 0, max_idx)
@@ -117,19 +116,10 @@ class RVEA(Algorithm):
             merged_pop, merged_fitness, v, (current_gen / self.max_gen) ** self.alpha
         )
 
-        def rv_adaptation(pop_obj, v):
-            v_temp = v * jnp.tile(
-                (jnp.nanmax(pop_obj, axis=0) - jnp.nanmin(pop_obj, axis=0)), (len(v), 1)
-            )
-
-            next_v = v_temp / jnp.tile(
-                jnp.sqrt(jnp.sum(v_temp**2, axis=1)).reshape(len(v), 1),
-                (1, jnp.shape(v)[1]),
-            )
-
-            return next_v
+        def rv_adaptation(pop_obj, v, v0):
+            return v0 * (jnp.nanmax(pop_obj, axis=0) - jnp.nanmin(pop_obj, axis=0))
 
-        def no_update(_pop_obj, v):
+        def no_update(_pop_obj, v, v0):
             return v
 
         v = jax.lax.cond(
@@ -137,6 +127,7 @@ class RVEA(Algorithm):
             rv_adaptation,
             no_update,
             survivor_fitness,
+            v,
             state.init_v,
         )
 
diff --git a/src/evox/algorithms/mo/rveaa.py b/src/evox/algorithms/mo/rveaa.py
new file mode 100644
index 00000000..2d106e51
--- /dev/null
+++ b/src/evox/algorithms/mo/rveaa.py
@@ -0,0 +1,215 @@
+# --------------------------------------------------------------------------------------
+# 1. RVEA algorithm is described in the following papers:
+#
+# Title: A Reference Vector Guided Evolutionary Algorithm for Many-Objective Optimization
+# Link: https://ieeexplore.ieee.org/document/7386636
+# --------------------------------------------------------------------------------------
+
+import jax
+import jax.numpy as jnp
+
+from evox.operators import mutation, crossover, selection
+from evox.operators.sampling import UniformSampling
+from evox import Algorithm, State, jit_class
+from evox.utils import cos_dist
+from evox.operators import non_dominated_sort
+
+
+@jax.jit
+def rv_regeneration(pop_obj, v, key):
+    """
+    Regenerate reference vectors regenerate strategy.
+    """
+    pop_obj = pop_obj - jnp.nanmin(pop_obj, axis=0)
+    cosine = cos_dist(pop_obj, v)
+
+    associate = jnp.nanargmax(cosine, axis=1)
+
+    invalid = jnp.sum(associate[:, jnp.newaxis] == jnp.arange(v.shape[0]), axis=0)
+    rand = jax.random.uniform(key, (v.shape[0], v.shape[1])) * jnp.nanmax(
+        pop_obj, axis=0
+    )
+    v = jnp.where(invalid[:, jnp.newaxis] == 0, rand, v)
+
+    return v
+
+
+@jax.jit
+def batch_truncation(pop, obj):
+    """
+    Use the batch truncation operator to select the best n solutions.
+    """
+    n = jnp.shape(pop)[0] // 2
+    cosine = cos_dist(obj, obj)
+    not_all_nan_rows = ~jnp.isnan(cosine).all(axis=1)
+    mask = jnp.eye(jnp.shape(cosine)[0], dtype=bool) & not_all_nan_rows[:, None]
+    cosine = jnp.where(mask, 0, cosine)
+
+    sorted_indices = jnp.sort(-cosine, axis=1)
+    rank = jnp.argsort(
+        jnp.where(jnp.isnan(sorted_indices[:, 0]), -jnp.inf, sorted_indices[:, 0])
+    )
+
+    mask = jnp.ones(jnp.shape(rank)[0], dtype=bool)
+    mask = mask.at[rank[:n]].set(False)[:, jnp.newaxis]
+
+    new_pop = jnp.where(mask, pop, jnp.nan)
+    new_obj = jnp.where(mask, obj, jnp.nan)
+
+    return new_pop, new_obj
+
+
+@jit_class
+class RVEAa(Algorithm):
+    """RVEAa algorithms (RVEA embedded with the reference vector regeneration strategy)
+
+    link: https://ieeexplore.ieee.org/document/7386636
+
+    Args:
+        alpha : The parameter controlling the rate of change of penalty. Defaults to 2.
+        fr : The frequency of reference vector adaptation. Defaults to 0.1.
+        max_gen : The maximum number of generations. Defaults to 100.
+        If the number of iterations is not 100, change the value based on the actual value.
+    """
+
+    def __init__(
+        self,
+        lb,
+        ub,
+        n_objs,
+        pop_size,
+        alpha=2,
+        fr=0.1,
+        max_gen=100,
+        selection_op=None,
+        mutation_op=None,
+        crossover_op=None,
+    ):
+        self.lb = lb
+        self.ub = ub
+        self.n_objs = n_objs
+        self.dim = lb.shape[0]
+        self.pop_size = pop_size
+        self.alpha = alpha
+        self.fr = fr
+        self.max_gen = max_gen
+
+        self.selection = selection_op
+        self.mutation = mutation_op
+        self.crossover = crossover_op
+
+        if self.selection is None:
+            self.selection = selection.ReferenceVectorGuided()
+        if self.mutation is None:
+            self.mutation = mutation.Polynomial((lb, ub))
+        if self.crossover is None:
+            self.crossover = crossover.SimulatedBinary()
+
+        self.sampling = UniformSampling(self.pop_size, self.n_objs)
+
+    def setup(self, key):
+        key, subkey1, subkey2, subkey3 = jax.random.split(key, 4)
+
+        v = self.sampling(subkey1)[0]
+        v0 = v
+        self.pop_size = v.shape[0]
+
+        population0 = (
+            jax.random.uniform(subkey2, shape=(self.pop_size, self.dim))
+            * (self.ub - self.lb)
+            + self.lb
+        )
+        population = jnp.concatenate(
+            [
+                population0,
+                jnp.full(shape=(self.pop_size, self.dim), fill_value=jnp.nan),
+            ],
+            axis=0,
+        )
+        v = jnp.concatenate(
+            [v, jax.random.uniform(subkey3, shape=(self.pop_size, self.n_objs))], axis=0
+        )
+
+        return State(
+            population=population,
+            fitness=jnp.zeros((self.pop_size * 2, self.n_objs)),
+            next_generation=population0,
+            reference_vector=v,
+            init_v=v0,
+            key=key,
+            gen=0,
+        )
+
+    def init_ask(self, state):
+        return state.population, state
+
+    def init_tell(self, state, fitness):
+        state = state.update(fitness=fitness)
+        return state
+
+    def ask(self, state):
+        key, subkey, x_key, mut_key = jax.random.split(state.key, 4)
+
+        population = state.population
+        no_nan_pop = ~jnp.isnan(population).all(axis=1)
+        max_idx = jnp.sum(no_nan_pop).astype(int)
+        pop = population[jnp.where(no_nan_pop, size=self.pop_size, fill_value=-1)]
+
+        mating_pool = jax.random.randint(subkey, (self.pop_size,), 0, max_idx)
+        crossovered = self.crossover(x_key, pop[mating_pool])
+        next_generation = self.mutation(mut_key, crossovered)
+        next_generation = jnp.clip(next_generation, self.lb, self.ub)
+
+        return next_generation, state.update(next_generation=next_generation, key=key)
+
+    def tell(self, state, fitness):
+        key, subkey = jax.random.split(state.key, 2)
+        current_gen = state.gen + 1
+
+        v = state.reference_vector
+        merged_pop = jnp.concatenate([state.population, state.next_generation], axis=0)
+
+        merged_fitness = jnp.concatenate([state.fitness, fitness], axis=0)
+
+        rank = non_dominated_sort(merged_fitness)
+        merged_fitness = jnp.where(rank[:, jnp.newaxis] == 0, merged_fitness, jnp.nan)
+        merged_pop = jnp.where(rank[:, jnp.newaxis] == 0, merged_pop, jnp.nan)
+
+        survivor, survivor_fitness = self.selection(
+            merged_pop, merged_fitness, v, (current_gen / self.max_gen) ** self.alpha
+        )
+
+        def rv_adaptation(pop_obj, v, v0):
+            return v0 * (jnp.nanmax(pop_obj, axis=0) - jnp.nanmin(pop_obj, axis=0))
+
+        def no_update(_pop_obj, v, v0):
+            return v
+
+        v_adapt = jax.lax.cond(
+            current_gen % (1 / self.fr) == 0,
+            rv_adaptation,
+            no_update,
+            survivor_fitness,
+            v[: self.pop_size],
+            state.init_v,
+        )
+
+        v_regen = rv_regeneration(survivor_fitness, v[self.pop_size :], subkey)
+        v = jnp.concatenate([v_adapt, v_regen], axis=0)
+
+        survivor, survivor_fitness = jax.lax.cond(
+            current_gen + 1 == self.max_gen,
+            batch_truncation,
+            lambda x, y: (x, y),
+            survivor,
+            survivor_fitness,
+        )
+
+        state = state.update(
+            population=survivor,
+            fitness=survivor_fitness,
+            reference_vector=v,
+            gen=current_gen,
+            key=key,
+        )
+        return state
diff --git a/src/evox/problems/numerical/cec2022_so.py b/src/evox/problems/numerical/cec2022_so.py
index 277352b6..f63eb69e 100644
--- a/src/evox/problems/numerical/cec2022_so.py
+++ b/src/evox/problems/numerical/cec2022_so.py
@@ -127,7 +127,8 @@ class OperatFunc:
         term1 = 1 / jnp.sqrt(diff_square_sum)
         term2 = jnp.exp(-0.5 * diff_square_sum / (sigma**2 * D))
         W = term1 * term2
-        W_norm = W / jnp.sum(W)
+        term1_nan = jnp.isnan(term1)
+        W_norm = lax.select(jnp.any(term1_nan), term1_nan / jnp.count_nonzero(term1_nan), W / jnp.sum(W))
 
         f = jnp.sum(W_norm * (lamb * fs + bias))
 


Test patch:
diff --git a/tests/test_multi_objective_algorithms.py b/tests/test_multi_objective_algorithms.py
index d69c6c2f..0ca250cf 100644
--- a/tests/test_multi_objective_algorithms.py
+++ b/tests/test_multi_objective_algorithms.py
@@ -200,6 +200,7 @@ def test_lmocso():
     )
     run_moea(algorithm)
 
+
 def test_im_moea():
     algorithm = algorithms.IMMOEA(
         lb=jnp.full(shape=(N,), fill_value=LB),
@@ -208,3 +209,13 @@ def test_im_moea():
         pop_size=POP_SIZE,
     )
     run_moea(algorithm)
+
+
+def test_rveaa():
+    algorithm = algorithms.RVEAa(
+        lb=jnp.full(shape=(N,), fill_value=LB),
+        ub=jnp.full(shape=(N,), fill_value=UB),
+        n_objs=M,
+        pop_size=POP_SIZE,
+    )
+    run_moea(algorithm)