Merge pull request #310 from electronsandstuff/pierce/crowded-comparison-argsort-issues

NSGA2 Generator NaN Handling in `crowded_comparison_argsort` Bugfix

Author: roussel-ryan

xopt/generators/ga/nsga2.py

@@ -1,6 +1,6 @@
 import numpy as np
 from pydantic import Field, Discriminator
-from typing import Dict, List, Optional, Annotated, Union
+from typing import Annotated
 import pandas as pd
 import os
 from datetime import datetime
@@ -25,6 +25,19 @@
 ########################################################################################################################
 
 
+def vocs_data_to_arr(data: list | np.ndarray) -> np.ndarray:
+    """Force data coming from VOCS object into 2D numpy array (or None) for compatibility with helper functions"""
+    if isinstance(data, list):
+        data = np.ndarray(list)
+    if data.size == 0:
+        return None
+    if len(data.shape) == 1:
+        return data[:, None]
+    if len(data.shape) == 2:
+        return data
+    raise ValueError(f"Unrecognized shape from VOCS data: {data.shape}")
+
+
 def get_crowding_distance(pop_f: np.ndarray) -> np.ndarray:
     """
     Calculates NSGA-II style crowding distance as described in [1].
@@ -66,16 +79,20 @@ def get_crowding_distance(pop_f: np.ndarray) -> np.ndarray:
 
 
 def crowded_comparison_argsort(
-    pop_f: np.ndarray, pop_g: Optional[np.ndarray] = None
+    pop_f: np.ndarray, pop_g: np.ndarray | None = None
 ) -> np.ndarray:
     """
     Sorts the objective functions by domination rank and then by crowding distance (crowded comparison operator).
-    Indices to individuals are returned in order of increasing value by crowded comparison operator.
+    Indices to individuals are returned in order of increasing value of fitness by crowded comparison operator.
+    That is, the least fit individuals are returned first.
+
+    Notes: NaN values are removed from the comparison and added back at the beginning (least fit direction) of
+    the sorted indices.
 
     Parameters
     ----------
     pop_f : np.ndarray
-        (M, N) numpy array where N is the number of individuals and M is the number of objectives
+        (N, M) numpy array where N is the number of individuals and M is the number of objectives
     pop_g : np.ndarray, optional
         The constraints, by default None
 
@@ -84,51 +101,69 @@ def crowded_comparison_argsort(
     np.ndarray
         Numpy array of indices to sorted individuals
     """
-    # Deal with NaNs
-    pop_f = np.copy(pop_f)
-    pop_f[~np.isfinite(pop_g)] = 1e300
+    # Check for non-finite values in both pop_f and pop_g
+    has_nan = np.any(~np.isfinite(pop_f), axis=1)
     if pop_g is not None:
-        pop_g = np.copy(pop_g)
-        pop_g[~np.isfinite(pop_g)] = 1e300
+        has_nan = has_nan | np.any(~np.isfinite(pop_g), axis=1)
+    nan_indices = np.where(has_nan)[0]
+    finite_indices = np.where(~has_nan)[0]
 
-    ranks = fast_dominated_argsort(pop_f, pop_g)
-    inds = []
+    # If all values are non-finite, return the original indices
+    if len(finite_indices) == 0:
+        return np.arange(pop_f.shape[0])
+
+    # Extract only finite values for processing
+    pop_f_finite = pop_f[finite_indices, :]
+
+    # Handle constraints if provided
+    pop_g_finite = None
+    if pop_g is not None:
+        pop_g_finite = pop_g[finite_indices, :]
+
+    # Apply domination ranking
+    ranks = fast_dominated_argsort(pop_f_finite, pop_g_finite)
+
+    # Calculate crowding distance and sort within each rank
+    sorted_finite_indices = []
     for rank in ranks:
-        dist = get_crowding_distance(pop_f[rank, :])
-        inds.extend(np.array(rank)[np.argsort(dist)[::-1]])
+        dist = get_crowding_distance(pop_f_finite[rank, :])
+        sorted_rank = np.array(rank)[np.argsort(dist)[::-1]]
+        sorted_finite_indices.extend(sorted_rank)
+
+    # Map back to original indices and put nans at end
+    sorted_original_indices = finite_indices[sorted_finite_indices]
+    final_sorted_indices = np.concatenate([sorted_original_indices, nan_indices])
 
-    return np.array(inds)[::-1]
+    return final_sorted_indices[::-1]
 
 
-def get_fitness(pop_f: np.ndarray, pop_g: np.ndarray) -> np.ndarray:
+def get_fitness(pop_f: np.ndarray, pop_g: np.ndarray | None = None) -> np.ndarray:
     """
     Get the "fitness" of each individual according to domination and crowding distance.
 
     Parameters
     ----------
     pop_f : np.ndarray
         The objectives
-    pop_g : np.ndarray
-        The constraints
+    pop_g : np.ndarray / None
+        The constraints, or None of no constraints
 
     Returns
     -------
     np.ndarray
         The fitness of each individual
     """
-    sort_ind = crowded_comparison_argsort(pop_f, pop_g)
-    fitness = np.argsort(sort_ind)
-    return fitness
+    return np.argsort(crowded_comparison_argsort(pop_f, pop_g))
 
 
 def generate_child_binary_tournament(
     pop_x: np.ndarray,
     pop_f: np.ndarray,
-    pop_g: np.ndarray,
+    pop_g: np.ndarray | None,
     bounds: np.ndarray,
     mutate: MutationOperator,
     crossover: CrossoverOperator,
-    fitness: Optional[np.ndarray] = None,
+    fitness: np.ndarray | None = None,
 ) -> np.ndarray:
     """
     Creates a single child from the population using binary tournament selection, crossover, and mutation.
@@ -143,8 +178,9 @@ def generate_child_binary_tournament(
         Decision variables of the population, shape (n_individuals, n_variables).
     pop_f : numpy.ndarray
         Objective function values of the population, shape (n_individuals, n_objectives).
-    pop_g : numpy.ndarray
+    pop_g : numpy.ndarray / None
         Constraint violation values of the population, shape (n_individuals, n_constraints).
+        None if no constraints.
     bounds : numpy.ndarray
         Bounds for decision variables, shape (2, n_variables) where bounds[0] are lower bounds
         and bounds[1] are upper bounds.
@@ -186,7 +222,7 @@ def generate_child_binary_tournament(
 
 
 def cull_population(
-    pop_x: np.ndarray, pop_f: np.ndarray, pop_g: np.ndarray, population_size: int
+    pop_x: np.ndarray, pop_f: np.ndarray, pop_g: np.ndarray | None, population_size: int
 ) -> np.ndarray:
     """
     Reduce population size by selecting the best individuals based on crowded comparison.
@@ -198,8 +234,8 @@ def cull_population(
     ----------
     pop_x : numpy.ndarray
         Decision variables of the population, shape (n_individuals, n_variables).
-    pop_f : numpy.ndarray
-        Objective function values of the population, shape (n_individuals, n_objectives).
+    pop_f : numpy.ndarray / None
+        Objective function values of the population, shape (n_individuals, n_objectives), None if no constraints.
     pop_g : numpy.ndarray
         Constraint violation values of the population, shape (n_individuals, n_constraints).
     population_size : int
@@ -210,9 +246,7 @@ def cull_population(
     numpy.ndarray
         Indices of selected individuals, shape (population_size,).
     """
-    inds = crowded_comparison_argsort(pop_f, pop_g)[::-1]
-    inds = inds[:population_size]
-    return inds
+    return crowded_comparison_argsort(pop_f, pop_g)[-population_size:]
 
 
 ########################################################################################################################
@@ -278,20 +312,20 @@ class NSGA2Generator(DeduplicatedGeneratorBase, StateOwner):
 
     population_size: int = Field(50, description="Population size")
     crossover_operator: Annotated[
-        Union[
-            SimulatedBinaryCrossover, DummyCrossover
-        ],  # Dummy placeholder to keep discriminator code from failing
+        (
+            SimulatedBinaryCrossover | DummyCrossover
+        ),  # Dummy placeholder to keep discriminator code from failing
         Discriminator("name"),
     ] = SimulatedBinaryCrossover()
     mutation_operator: Annotated[
-        Union[
-            PolynomialMutation, DummyMutation
-        ],  # Dummy placeholder to keep discriminator code from failing
+        (
+            PolynomialMutation | DummyMutation
+        ),  # Dummy placeholder to keep discriminator code from failing
         Discriminator("name"),
     ] = PolynomialMutation()
 
     # Output options
-    output_dir: Optional[str] = None
+    output_dir: str | None = None
     checkpoint_freq: int = Field(
         -1,
         description="How often (in generations) to save checkpoints (set to -1 to disable)",
@@ -302,7 +336,7 @@ class NSGA2Generator(DeduplicatedGeneratorBase, StateOwner):
     _output_dir_setup: bool = (
         False  # Used in initializing the directory. PLEASE DO NOT CHANGE
     )
-    _logger: Optional[logging.Logger] = None
+    _logger: logging.Logger | None = None
 
     # Metadata
     fevals: int = Field(
@@ -315,7 +349,7 @@ class NSGA2Generator(DeduplicatedGeneratorBase, StateOwner):
     n_candidates: int = Field(
         0, description="The number of candidate solutions generated so far"
     )
-    history_idx: List[List[int]] = Field(
+    history_idx: list[list[int]] = Field(
         default=[],
         description="Xopt indices of the individuals in each population",
     )
@@ -326,15 +360,15 @@ class NSGA2Generator(DeduplicatedGeneratorBase, StateOwner):
     )
 
     # The population and returned children
-    pop: List[Dict] = Field(default=[])
-    child: List[Dict] = Field(default=[])
+    pop: list[dict] = Field(default=[])
+    child: list[dict] = Field(default=[])
 
     def model_post_init(self, context):
         # Get a unique logger per object
         self._logger = logging.getLogger(f"{__name__}.NSGA2Generator.{id(self)}")
         self._logger.setLevel(self.log_level)
 
-    def _generate(self, n_candidates: int) -> List[Dict]:
+    def _generate(self, n_candidates: int) -> list[dict]:
         self.ensure_output_dir_setup()
         start_t = time.perf_counter()
 
@@ -347,7 +381,7 @@ def _generate(self, n_candidates: int) -> List[Dict]:
             candidates = []
             pop_x = self.vocs.variable_data(self.pop).to_numpy()
             pop_f = self.vocs.objective_data(self.pop).to_numpy()
-            pop_g = self.vocs.constraint_data(self.pop).to_numpy()
+            pop_g = vocs_data_to_arr(self.vocs.constraint_data(self.pop).to_numpy())
             fitness = get_fitness(pop_f, pop_g)
             for _ in range(n_candidates):
                 candidates.append(
@@ -419,7 +453,7 @@ def add_data(self, new_data: pd.DataFrame):
             idx = cull_population(
                 self.vocs.variable_data(self.pop).to_numpy(),
                 self.vocs.objective_data(self.pop).to_numpy(),
-                self.vocs.constraint_data(self.pop).to_numpy(),
+                vocs_data_to_arr(self.vocs.constraint_data(self.pop).to_numpy()),
                 self.population_size,
             )
             self.pop = [self.pop[i] for i in idx]

xopt/generators/utils.py

@@ -1,9 +1,8 @@
 # from xopt.generator import Generator
 import numpy as np
-from typing import List, Optional
 
 
-def get_domination(pop_f: np.ndarray, pop_g: Optional[np.ndarray] = None) -> np.ndarray:
+def get_domination(pop_f: np.ndarray, pop_g: np.ndarray | None = None) -> np.ndarray:
     """
     Compute domination matrix for a population based on objective values and constraints. Determines domination
     relationships between all pairs of individuals in a population.
@@ -46,7 +45,7 @@ def get_domination(pop_f: np.ndarray, pop_g: Optional[np.ndarray] = None) -> np.
     return dom
 
 
-def fast_dominated_argsort_internal(dom: np.ndarray) -> List[np.ndarray]:
+def fast_dominated_argsort_internal(dom: np.ndarray) -> list[np.ndarray]:
     """
     Used inside of `fast_dominated_argsort`. Call that function instead.
 
@@ -83,8 +82,8 @@ def fast_dominated_argsort_internal(dom: np.ndarray) -> List[np.ndarray]:
 
 
 def fast_dominated_argsort(
-    pop_f: np.ndarray, pop_g: Optional[np.ndarray] = None
-) -> List[np.ndarray]:
+    pop_f: np.ndarray, pop_g: np.ndarray | None = None
+) -> list[np.ndarray]:
     """
     Performs a dominated sort on matrix of objective function values O.  This is a numpy implementation of the algorithm
     described in [1].

xopt/tests/generators/ga/test_nsga2.py

@@ -5,12 +5,14 @@
 
 import numpy as np
 import pandas as pd
+import pytest
 
 from xopt.base import Xopt
 from xopt.evaluator import Evaluator
 from xopt.generators.ga.nsga2 import (
     NSGA2Generator,
     generate_child_binary_tournament,
+    crowded_comparison_argsort,
 )
 from xopt.generators.ga.operators import PolynomialMutation, SimulatedBinaryCrossover
 from xopt.resources.test_functions.tnk import evaluate_TNK, tnk_vocs
@@ -397,3 +399,80 @@ def compare(val_a, val_b):
             np.random.seed(42)
             samples = X.generator.generate(1)
             X.evaluate_data(samples)
+
+
+@pytest.mark.parametrize(
+    "pop_f, pop_g, expected_indices_options",
+    [
+        # Two individuals in different ranks
+        (np.array([[1.0, 2.0], [2.0, 3.0]]), None, [np.array([1, 0])]),
+        # Non-dominated, different crowding distances
+        (
+            np.array([[1.0, 3.0], [2.0, 2.0], [3.0, 1.0]]),
+            None,
+            [np.array([1, 2, 0]), np.array([1, 0, 2])],
+        ),
+        # NaN values
+        (
+            np.array([[1.0, 2.0], [np.nan, 3.0], [2.0, 1.0]]),
+            None,
+            [np.array([1, 2, 0]), np.array([1, 0, 2])],
+        ),
+        # Constrained
+        (
+            np.array([[2.0, 2.0], [1.0, 1.0]]),
+            np.array([[-1.0, -1.0], [1.0, -1.0]]),
+            [np.array([1, 0])],
+        ),
+        # Multiple individuals with same rank but potentially same crowding distances
+        (
+            np.array([[1.0, 3.0], [2.0, 2.0], [3.0, 1.0], [1.5, 2.5]]),
+            None,
+            [
+                np.array([3, 1, 2, 0]),
+                np.array([1, 3, 2, 0]),
+                np.array([3, 1, 0, 2]),
+                np.array([1, 3, 0, 2]),
+            ],
+        ),
+        # NaN values and constraints
+        (
+            np.array([[1.0, 2.0], [np.nan, 3.0], [2.0, 1.0]]),
+            np.array([[-1.0, -1.0, -1.0], [0.0, 0.0, 0.0], [1.0, 0.0, 0.0]]),
+            [np.array([1, 2, 0])],
+        ),
+        # All NaN
+        (
+            np.array([[np.nan, np.nan], [np.nan, np.nan]]),
+            None,
+            [np.array([0, 1]), np.array([1, 0])],
+        ),
+    ],
+)
+def test_crowded_comparison_argsort(pop_f, pop_g, expected_indices_options):
+    """
+    Test the crowded_comparison_argsort function with various explicit input.
+
+    Parameters
+    ----------
+    pop_f : numpy.ndarray
+        Objective values
+    pop_g : numpy.ndarray or None
+        Constraint values
+    expected_indices_options : list of numpy.ndarray
+        List of valid expected sorted indices
+    """
+    # Call the function
+    result = crowded_comparison_argsort(pop_f, pop_g)
+
+    # Check if the result matches any of the expected options
+    matches_any = any(
+        np.array_equal(result, expected) for expected in expected_indices_options
+    )
+
+    if not matches_any:
+        message = (
+            f"Result {result} doesn't match any expected ordering.\n"
+            f"Expected one of: {expected_indices_options}"
+        )
+        assert False, message