From 7755719477b5893544943f25107b57de26782e55 Mon Sep 17 00:00:00 2001
From: wmartens <waldemar.martens@esa.int>
Date: Thu, 30 Jan 2025 15:50:36 +0000
Subject: [PATCH 01/10] added KhanFunction class

---
 pyoptgra/optgra.py | 93 ++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 93 insertions(+)

diff --git a/pyoptgra/optgra.py b/pyoptgra/optgra.py
index 02282f5..cf40ab2 100644
--- a/pyoptgra/optgra.py
+++ b/pyoptgra/optgra.py
@@ -15,6 +15,7 @@
 from collections import deque
 from math import isfinite
 from typing import List, Tuple, Union
+import numpy as np
 
 from pygmo import s_policy, select_best
 
@@ -27,6 +28,98 @@
 )
 
 
+class KhanFunction:
+    """Function to smothly enforce optimisation parameter bounds as Michal Khan used to do:
+
+    .. math::
+
+        x = \frac{x_{max} + x_{min}}{2} + \frac{x_{max} - x_{min}}{2} \cdot \sin(x_{optgra})
+
+    Where :math:`x` is the pagmo decision vector and :math:`x_{optgra}` is the decision vector
+    passed to OPTGRA. In this way parameter bounds are guaranteed to be satisfied, but the gradients
+    near the bounds approach zero.
+    """
+
+    def __init__(self, lb: List[float], ub: List[float]):
+        """Constructor
+
+        Parameters
+        ----------
+        lb : List[float]
+            Lower pagmo parameter bounds
+        ub : List[float]
+            Upper pagmo parameter bounds
+        """
+        self._lb = lb
+        self._ub = ub
+
+    def eval(self, x_optgra: List[float]) -> List[float]:
+        """Convert :math:`x_{optgra}` to :math:`x`.
+
+        Parameters
+        ----------
+        x_optgra : List[float]
+            Decision vector passed to OPTGRA
+
+        Returns
+        -------
+        List[float]
+            Pagmo decision vector
+        """
+        return (self._ub + self._lb) / 2 + (self._ub - self._lb) / 2 * np.sin(x_optgra)
+
+    def eval_inv(self, x: List[float]) -> List[float]:
+        """Convert :math:`x` to :math:`x_{optgra}`.
+
+        Parameters
+        ----------
+        x : List[float]
+            Pagmo decision vector
+
+        Returns
+        -------
+        List[float]
+            Decision vector passed to OPTGRA
+
+        """
+        arg = (2 * x - self._ub - self._lb) / (self._ub - self._lb)
+        # TODO: check that arg is in the interval [-1, 1]
+        return np.asin(arg)
+
+    def eval_grad(self, x_optgra: List[float]) -> List[float]:
+        """Gradient of ``eval`` function.
+
+        Parameters
+        ----------
+        x_optgra : List[float]
+            Decision vector passed to OPTGRA
+
+        Returns
+        -------
+        List[float]
+            Pagmo decision vector
+        """
+        return (self._ub - self._lb) / 2 * np.cos(x_optgra)
+
+    def eval_grad_inv(self, x_optgra: List[float]) -> List[float]:
+        """Gradient of ``eval_inv`` method.
+
+        Parameters
+        ----------
+        x_optgra : List[float]
+            Decision vector passed to OPTGRA
+
+        Returns
+        -------
+        List[float]
+            Pagmo decision vector
+        """
+        return -1 / (
+            (self._lb - self._ub)
+            * np.sqrt(((self._lb - x_optgra) * (x_optgra - self._ub)) / (self._ub - self._lb) ** 2)
+        )
+
+
 class optgra:
     """
     This class is a user defined algorithm (UDA) providing a wrapper around OPTGRA, which is written

From ec0c117a7736dea6d59dae54cca53a79f4485c37 Mon Sep 17 00:00:00 2001
From: wmartens <waldemar.martens@esa.int>
Date: Thu, 30 Jan 2025 16:59:56 +0000
Subject: [PATCH 02/10] finished Khan function

---
 pyoptgra/optgra.py   | 200 ++++++++++++++++++++++++++++++-------------
 tests/python/test.py |  87 ++++++++++++++-----
 2 files changed, 207 insertions(+), 80 deletions(-)

diff --git a/pyoptgra/optgra.py b/pyoptgra/optgra.py
index cf40ab2..cc142c8 100644
--- a/pyoptgra/optgra.py
+++ b/pyoptgra/optgra.py
@@ -14,7 +14,7 @@
 
 from collections import deque
 from math import isfinite
-from typing import List, Tuple, Union
+from typing import List, Tuple, Union, Callable
 import numpy as np
 
 from pygmo import s_policy, select_best
@@ -40,6 +40,7 @@ class KhanFunction:
     near the bounds approach zero.
     """
 
+    # TODO: coversion function shall only be called on parameters with finite bounds
     def __init__(self, lb: List[float], ub: List[float]):
         """Constructor
 
@@ -53,6 +54,53 @@ def __init__(self, lb: List[float], ub: List[float]):
         self._lb = lb
         self._ub = ub
 
+        # determine finite lower/upper bounds
+        finite_lb = np.isfinite(lb)
+        finite_ub = np.isfinite(ub)
+
+        # we only support cases where both lower and upper bounds are finite if given
+        check = np.where(finite_lb != finite_ub)[0]
+        if any(check):
+            raise ValueError(
+                "When using Khan bounds, both lower and upper bound for bounded parameters "
+                "must be finite."
+                f"Detected mismatch at decision vector indices: {check}"
+            )
+        # store the mask of finite bounds
+        self.mask = finite_ub
+        self._lb_masked = lb[self.mask]
+        self._ub_masked = ub[self.mask]
+
+    def _eval(self, x_optgra_masked: List[float]) -> List[float]:
+        return (self._ub_masked + self._lb_masked) / 2 + (
+            self._ub_masked - self._lb_masked
+        ) / 2 * np.sin(x_optgra_masked)
+
+    def _eval_inv(self, x_masked: List[float]) -> List[float]:
+        arg = (2 * x_masked - self._ub_masked - self._lb_masked) / (
+            self._ub_masked - self._lb_masked
+        )
+        # TODO: check that arg is in the interval [-1, 1]
+        return np.asin(arg)
+
+    def _eval_grad(self, x_optgra_masked: List[float]) -> List[float]:
+        return (self._ub_masked - self._lb_masked) / 2 * np.cos(x_optgra_masked)
+
+    def _eval_grad_inv(self, x_optgra_masked: List[float]) -> List[float]:
+        return -1 / (
+            (self._lb_masked - self._ub_masked)
+            * np.sqrt(
+                ((self._lb_masked - x_optgra_masked) * (x_optgra_masked - self._ub_masked))
+                / (self._ub_masked - self._lb_masked) ** 2
+            )
+        )
+
+    def _apply_to_subset(self, x: List[float], func: Callable):
+        """Apply a given function only to a subset of x defined by self.mask."""
+        result = x.copy()  # Create a copy to preserve the original structure
+        result[self.mask] = func(x[self.mask])  # Apply the function only to the selected subset
+        return result
+
     def eval(self, x_optgra: List[float]) -> List[float]:
         """Convert :math:`x_{optgra}` to :math:`x`.
 
@@ -66,7 +114,7 @@ def eval(self, x_optgra: List[float]) -> List[float]:
         List[float]
             Pagmo decision vector
         """
-        return (self._ub + self._lb) / 2 + (self._ub - self._lb) / 2 * np.sin(x_optgra)
+        return self._apply_to_subset(x_optgra, self._eval)
 
     def eval_inv(self, x: List[float]) -> List[float]:
         """Convert :math:`x` to :math:`x_{optgra}`.
@@ -82,9 +130,7 @@ def eval_inv(self, x: List[float]) -> List[float]:
             Decision vector passed to OPTGRA
 
         """
-        arg = (2 * x - self._ub - self._lb) / (self._ub - self._lb)
-        # TODO: check that arg is in the interval [-1, 1]
-        return np.asin(arg)
+        return self._apply_to_subset(x, self.eval_inv)
 
     def eval_grad(self, x_optgra: List[float]) -> List[float]:
         """Gradient of ``eval`` function.
@@ -99,7 +145,7 @@ def eval_grad(self, x_optgra: List[float]) -> List[float]:
         List[float]
             Pagmo decision vector
         """
-        return (self._ub - self._lb) / 2 * np.cos(x_optgra)
+        return self._apply_to_subset(x_optgra, self._eval_grad)
 
     def eval_grad_inv(self, x_optgra: List[float]) -> List[float]:
         """Gradient of ``eval_inv`` method.
@@ -114,10 +160,7 @@ def eval_grad_inv(self, x_optgra: List[float]) -> List[float]:
         List[float]
             Pagmo decision vector
         """
-        return -1 / (
-            (self._lb - self._ub)
-            * np.sqrt(((self._lb - x_optgra) * (x_optgra - self._ub)) / (self._ub - self._lb) ** 2)
-        )
+        return self._apply_to_subset(x_optgra, self._eval_grad_inv)
 
 
 class optgra:
@@ -159,18 +202,24 @@ def _wrap_fitness_func(
         problem,
         bounds_to_constraints: bool = True,
         force_bounds: bool = False,
+        khan_function: KhanFunction = None,
     ):
         def wrapped_fitness(x):
 
+            # if Khan function is used, we first need to convert to pagmo parameters
+            if khan_function:
+                x = KhanFunction.eval(x)
+
             fixed_x = x
             lb, ub = problem.get_bounds()
 
             if force_bounds:
-                for i in range(problem.get_nx()):
-                    if x[i] < lb[i]:
-                        fixed_x[i] = lb[i]
-                    if x[i] > ub[i]:
-                        fixed_x[i] = ub[i]
+                x = np.clip(x, lb, ub).tolist()
+                # for i in range(problem.get_nx()):
+                #     if x[i] < lb[i]:
+                #         fixed_x[i] = lb[i]
+                #     if x[i] > ub[i]:
+                #         fixed_x[i] = ub[i]
 
             result = deque(problem.fitness(fixed_x))
 
@@ -191,22 +240,33 @@ def wrapped_fitness(x):
         return wrapped_fitness
 
     @staticmethod
-    def _wrap_gradient_func(problem, bounds_to_constraints: bool = True, force_bounds=False):
+    def _wrap_gradient_func(
+        problem,
+        bounds_to_constraints: bool = True,
+        force_bounds=False,
+        khan_function: KhanFunction = None,
+    ):
 
         sparsity_pattern = problem.gradient_sparsity()
 
         shape = (problem.get_nf(), problem.get_nx())
 
         def wrapped_gradient(x):
+
+            # if Khan function is used, we first need to convert to pagmo parameters
+            if khan_function:
+                x = KhanFunction.eval(x)
+
             lb, ub = problem.get_bounds()
             fixed_x = x
 
             if force_bounds:
-                for i in range(problem.get_nx()):
-                    if x[i] < lb[i]:
-                        fixed_x[i] = lb[i]
-                    if x[i] > ub[i]:
-                        fixed_x[i] = ub[i]
+                np.clip(x, lb, ub).tolist()
+                # for i in range(problem.get_nx()):
+                #     if x[i] < lb[i]:
+                #         fixed_x[i] = lb[i]
+                #     if x[i] > ub[i]:
+                #         fixed_x[i] = ub[i]
 
             sparse_values = problem.gradient(fixed_x)
 
@@ -258,6 +318,7 @@ def __init__(
         merit_function_threshold: float = 1e-6,
         # bound_constraints_scalar: float = 1,
         force_bounds: bool = False,
+        khan_bounds: bool = False,
         optimization_method: int = 2,
         log_level: int = 0,
     ) -> None:
@@ -273,38 +334,49 @@ def __init__(
 
         Args:
 
-            max_iterations: maximum number of total iterations max_correction_iterations: number of
-            constraint correction iterations in the beginning
+            max_iterations: maximum number of total iterations
+            max_correction_iterations: number of
+                constraint correction iterations in the beginning
                 If no feasible solution is found within that many iterations, Optgra aborts
             max_distance_per_iteration: maximum scaled distance traveled in each iteration
-            perturbation_for_snd_order_derivatives: Used as delta for numerically computing second
-            order errors
-                of the constraints in the optimization step
+            perturbation_for_snd_order_derivatives: Used as delta for numerically computing
+                second order errors of the constraints in the optimization step
             variable_scaling_factors: optional - Scaling factors for the input variables.
                 If passed, must be positive and as many as there are variables
             variable_types: optional - Flags to set variables to either free (0) or fixed (1). Fixed
-            variables
-                are also called parameters in sensitivity analysis. If passed, must be as many flags
-                as there are variables
+                variables are also called parameters in sensitivity analysis. If passed, must be as
+                many flags as there are variables
             constraint_priorities: optional - lower constraint priorities are fulfilled earlier.
                 During the initial constraint correction phase, only constraints with a priority at
                 most k are considered in iteration k. Defaults to zero, so that all constraints are
                 considered from the beginning.
             bounds_to_constraints: optional - if true (default), translate box bounds of the given
-            problems into
-                inequality constraints for optgra. Note that when also passing constraint
-                priorities, the original constraints of the problem come first, followed by those
-                derived from the lower box bounds, then those from the upper box bounds. Infinite
-                bounds are ignored and not counted.
+                problems into inequality constraints for optgra. Note that when also passing
+                constraint priorities, the original constraints of the problem come first, followed
+                by those derived from the lower box bounds, then those from the upper box bounds.
+                Infinite bounds are ignored and not counted.
             bound_constraints_tolerance: optional - constraint tolerance for the constraints derived
-            from bounds merit_function_threshold: optional - convergence threshold for merit
-            function force_bounds: optional - whether to force the bounds given by the problem. If
-            false (default), the
-                fitness function might also be called with values of x that are outside of the
-                bounds. Set to true if the fitness function cannot handle that.
+                from bounds
+            merit_function_threshold: optional - convergence threshold for merit
+                function
+            force_bounds: optional - whether to force the bounds given by the problem. If
+                false (default), the fitness function might also be called with values of x that are
+                outside of the bounds. Set to true if the fitness function cannot handle that.
+                If active, the gradients evaluated near the bounds will be inacurate potentially
+                leading to convergence issues.
+            khan_bounds: optional - whether to gracefully enforce bounds on the decision vector
+                using Michael Khan's method:
+
+                .. math::
+
+                    x = \frac{x_{max} + x_{min}}{2} + \frac{x_{max} - x_{min}}{2} \cdot \sin(x_{optgra})
+
+                Where :math:`x` is the pagmo decision vector and :math:`x_{optgra}` is the decision
+                vector passed to OPTGRA. In this way parameter bounds are guaranteed to be
+                satisfied, but the gradients near the bounds approach zero. By default False.
             optimization_method: select 0 for steepest descent, 1 for modified spectral conjugate
-            gradient method,
-                2 for spectral conjugate gradient method and 3 for conjugate gradient method
+                gradient method, 2 for spectral conjugate gradient method and 3 for conjugate
+                gradient method
             log_level: Control the original screen output of OPTGRA. 0 has no output,
                 4 and higher have maximum output`. Set this to 0 if you want to use the pygmo
                 logging system based on `set_verbosity()`.
@@ -331,6 +403,7 @@ def __init__(
         self.merit_function_threshold = merit_function_threshold
 
         self.force_bounds = force_bounds
+        self.khan_bounds = khan_bounds
         # self.bound_violation_penalty = bound_violation_penalty
 
         self.log_level = log_level
@@ -481,14 +554,17 @@ def evolve(self, population):
 
         idx = list(population.get_ID()).index(selected[0][0])
 
+        # optional Khan function to enforce parameter bounds
+        khan_function = KhanFunction(*problem.get_bounds()) if self.khan_bounds else None
+
         fitness_func = optgra._wrap_fitness_func(
-            problem, self.bounds_to_constraints, self.force_bounds
+            problem, self.bounds_to_constraints, self.force_bounds, khan_function
         )
         grad_func = None
         derivatives_computation = 2
         if problem.has_gradient():
             grad_func = optgra._wrap_gradient_func(
-                problem, self.bounds_to_constraints, self.force_bounds
+                problem, self.bounds_to_constraints, self.force_bounds, khan_function
             )
             derivatives_computation = 1
 
@@ -522,8 +598,10 @@ def evolve(self, population):
         if len(variable_types) == 0:
             variable_types = [0 for _ in range(problem.get_nx())]
 
+        # get initial x
+        x0 = population.get_x()[idx]
         result = optimize(
-            initial_x=population.get_x()[idx],
+            initial_x=khan_function.eval_inv(x0) if khan_function else x0,
             constraint_types=constraint_types,
             fitness_callback=fitness_func,
             gradient_callback=grad_func,
@@ -547,23 +625,25 @@ def evolve(self, population):
 
         best_x, best_f, finopt = result
 
-        violated = False
+        # if a Khan function is used we first need to convert to pagmo parameters
+        if khan_function:
+            best_x = khan_function.eval(best_x)
+
         if self.force_bounds:
             lb, ub = problem.get_bounds()
-            for i in range(problem.get_nx()):
-                if best_x[i] < lb[i]:
-                    best_x[i] = lb[i]
-                    violated = True
-                if best_x[i] > ub[i]:
-                    best_x[i] = ub[i]
-                    violated = True
-        if violated:
-            population.set_x(idx, best_x)
-        else:
-            # merit function is last, constraints are from 0 to problem.get_nc(),
-            # we ignore bound-derived constraints
-            # pagmo_fitness = [best_f[-1]] + best_f[0 : problem.get_nc()]
-            population.set_x(idx, best_x)  # , list(pagmo_fitness))
+            best_x = np.clip(best_x, lb, ub).tolist()
+            # for i in range(problem.get_nx()):
+            #     if best_x[i] < lb[i]:
+            #         best_x[i] = lb[i]
+            #         violated = True
+            #     if best_x[i] > ub[i]:
+            #         best_x[i] = ub[i]
+            #         violated = True
+
+        # merit function is last, constraints are from 0 to problem.get_nc(),
+        # we ignore bound-derived constraints
+        # pagmo_fitness = [best_f[-1]] + best_f[0 : problem.get_nc()]
+        population.set_x(idx, best_x)  # , list(pagmo_fitness))
 
         return population
 
@@ -822,6 +902,7 @@ def get_extra_info(self) -> str:
             + "bound_constraints_tolerance = {bound_constraints_tolerance}, "
             + "merit_function_threshold = {merit_function_threshold}, "
             + "force_bounds = {force_bounds}, "
+            + "khan_bounds = {khan_bounds}, "
             + "optimization_method = {optimization_method}, "
             + "log_level = {log_level}"
             + "verbosity = {verbosity}"
@@ -837,6 +918,7 @@ def get_extra_info(self) -> str:
             bound_constraints_tolerance=self.bound_constraints_tolerance,
             merit_function_threshold=self.merit_function_threshold,
             force_bounds=self.force_bounds,
+            khan_bounds=self.khan_bounds,
             optimization_method=self.optimization_method,
             log_level=self.log_level,
             verbosity=self.verbosity,
diff --git a/tests/python/test.py b/tests/python/test.py
index 5531734..3148446 100644
--- a/tests/python/test.py
+++ b/tests/python/test.py
@@ -74,12 +74,7 @@ def fitness(self, x):
             - x[1]
         )
         ci2 = -(
-            8 * x[5] * (x[5] ** 2 - x[4])
-            - 2 * (1 - x[5])
-            + x[4] ** 2
-            - x[3]
-            + x[3] ** 2
-            - x[4]
+            8 * x[5] * (x[5] ** 2 - x[4]) - 2 * (1 - x[5]) + x[4] ** 2 - x[3] + x[3] ** 2 - x[4]
         )
         return [obj, ce1, ce2, ce3, ce4, ci1, ci2]
 
@@ -189,9 +184,7 @@ def constructor_test(self):
 
         # Check that negative perturbation is rejected
         with self.assertRaises(ValueError):
-            _ = pygmo.algorithm(
-                pyoptgra.optgra(perturbation_for_snd_order_derivatives=-1)
-            )
+            _ = pygmo.algorithm(pyoptgra.optgra(perturbation_for_snd_order_derivatives=-1))
 
         # Valid constructor
         pygmo.algorithm(pyoptgra.optgra())
@@ -400,9 +393,7 @@ def archipelago_evolve_test(self):
         a.evolve(10)
         a2 = deepcopy(a)
         a.wait_check()
-        a = archipelago(
-            5, udi=mp_island(), algo=pyoptgra.optgra(), prob=rosenbrock(), pop_size=10
-        )
+        a = archipelago(5, udi=mp_island(), algo=pyoptgra.optgra(), prob=rosenbrock(), pop_size=10)
         a.evolve(10)
         a.evolve(10)
         str(a)
@@ -436,9 +427,7 @@ def archipelago_pickle_test(self):
 
         a = archipelago(5, algo=pyoptgra.optgra(), prob=rosenbrock(), pop_size=10)
         self.assertEqual(repr(a), repr(loads(dumps(a))))
-        a = archipelago(
-            5, algo=pyoptgra.optgra(), prob=_prob(), pop_size=10, udi=mp_island()
-        )
+        a = archipelago(5, algo=pyoptgra.optgra(), prob=_prob(), pop_size=10, udi=mp_island())
         self.assertEqual(repr(a), repr(loads(dumps(a))))
 
     def prepare_sensitivity_input_check_test(self):
@@ -520,9 +509,7 @@ def sensitivity_matrices_test(self):
 
         matrices = opt.sensitivity_matrices()
         self.assertEqual(len(matrices), 5)
-        self.assertEqual(
-            len(matrices[0]), 6 + 12
-        )  # luksan_vlcek has 6 constraints and 12 bounds
+        self.assertEqual(len(matrices[0]), 6 + 12)  # luksan_vlcek has 6 constraints and 12 bounds
         self.assertLessEqual(max(matrices[0]), 1)
         self.assertGreaterEqual(min(matrices[0]), 0)
 
@@ -650,9 +637,66 @@ def get_nic(self):
         self.assertFalse(extracted._bounds_violated)
 
         # check that population has valid members
-        algo = pygmo.algorithm(
-            pyoptgra.optgra(force_bounds=True, bounds_to_constraints=False)
-        )
+        algo = pygmo.algorithm(pyoptgra.optgra(force_bounds=True, bounds_to_constraints=False))
+        prob = pygmo.problem(_prob_bound_test_no_gradient())
+        pop = pygmo.population(prob, size=0)
+        pop.push_back([2.47192039, -1.45880516, -9.03600606, -9.33306356, 3.85509973])
+        extracted = pop.problem.extract(_prob_bound_test_no_gradient)
+        self.assertFalse(extracted._bounds_violated)
+        pop = algo.evolve(pop)
+        extracted = pop.problem.extract(_prob_bound_test_no_gradient)
+        self.assertFalse(extracted._bounds_violated)
+        lb, ub = prob.get_bounds()
+        for i in range(prob.get_nx()):
+            self.assertTrue(pop.champion_x[i] >= lb[i])
+            self.assertTrue(pop.champion_x[i] <= ub[i])
+
+    def khan_bounds_test(self):
+        class _prob_bound_test_no_gradient(object):
+            def __init__(self):
+                self._bounds_violated = False
+
+            def get_bounds(self):
+                return ([-10, -10, -10, -10, -10], [10, 10, 10, 10, 10])
+
+            def fitness(self, x):
+                result = [
+                    sum(x),
+                    sum([(x[i] - 5 * i + 10) ** 2 for i in range(len(x))]),
+                ]
+                lb, ub = self.get_bounds()
+                for i in range(len(lb)):
+                    if x[i] < lb[i]:
+                        self._bounds_violated = True
+                    if x[i] > ub[i]:
+                        self._bounds_violated = True
+                return result
+
+            def get_nic(self):
+                return 1
+
+        # check bounds violation with normal optgra
+        algo = pygmo.algorithm(pyoptgra.optgra(force_bounds=False))
+        prob = pygmo.problem(_prob_bound_test_no_gradient())
+        pop = pygmo.population(prob, size=1)
+        extracted = pop.problem.extract(_prob_bound_test_no_gradient)
+        self.assertFalse(extracted._bounds_violated)
+        pop = algo.evolve(pop)
+        extracted = pop.problem.extract(_prob_bound_test_no_gradient)
+        self.assertTrue(extracted._bounds_violated)
+
+        # check bounds are forced when setting the argument
+        algo = pygmo.algorithm(pyoptgra.optgra(force_bounds=True))
+        prob = pygmo.problem(_prob_bound_test())
+        pop = pygmo.population(prob, size=1)
+        extracted = pop.problem.extract(_prob_bound_test)
+        self.assertFalse(extracted._bounds_violated)
+        pop = algo.evolve(pop)
+        extracted = pop.problem.extract(_prob_bound_test)
+        self.assertFalse(extracted._bounds_violated)
+
+        # check that population has valid members
+        algo = pygmo.algorithm(pyoptgra.optgra(force_bounds=True, bounds_to_constraints=False))
         prob = pygmo.problem(_prob_bound_test_no_gradient())
         pop = pygmo.population(prob, size=0)
         pop.push_back([2.47192039, -1.45880516, -9.03600606, -9.33306356, 3.85509973])
@@ -679,5 +723,6 @@ def verbosity_test(self):
         with self.assertRaises(ValueError):
             algo.set_verbosity(1)
 
+
 if __name__ == "__main__":
     unittest.main()

From 6acf43eb78711ac55f296c0e72200243dbe59330 Mon Sep 17 00:00:00 2001
From: wmartens <waldemar.martens@esa.int>
Date: Fri, 31 Jan 2025 11:45:27 +0000
Subject: [PATCH 03/10] finished khan_function and vectorized wrappers

---
 pyoptgra/optgra.py   | 181 +++++++++++++++++++++++--------------------
 tests/python/test.py |  49 ++++++------
 2 files changed, 120 insertions(+), 110 deletions(-)

diff --git a/pyoptgra/optgra.py b/pyoptgra/optgra.py
index cc142c8..7d1cf8b 100644
--- a/pyoptgra/optgra.py
+++ b/pyoptgra/optgra.py
@@ -40,7 +40,6 @@ class KhanFunction:
     near the bounds approach zero.
     """
 
-    # TODO: coversion function shall only be called on parameters with finite bounds
     def __init__(self, lb: List[float], ub: List[float]):
         """Constructor
 
@@ -51,12 +50,12 @@ def __init__(self, lb: List[float], ub: List[float]):
         ub : List[float]
             Upper pagmo parameter bounds
         """
-        self._lb = lb
-        self._ub = ub
+        self._lb = np.asarray(lb)
+        self._ub = np.asarray(ub)
 
         # determine finite lower/upper bounds
-        finite_lb = np.isfinite(lb)
-        finite_ub = np.isfinite(ub)
+        finite_lb = np.isfinite(self._lb)
+        finite_ub = np.isfinite(self._ub)
 
         # we only support cases where both lower and upper bounds are finite if given
         check = np.where(finite_lb != finite_ub)[0]
@@ -71,22 +70,27 @@ def __init__(self, lb: List[float], ub: List[float]):
         self._lb_masked = lb[self.mask]
         self._ub_masked = ub[self.mask]
 
-    def _eval(self, x_optgra_masked: List[float]) -> List[float]:
+    def _eval(self, x_optgra_masked: np.ndarray) -> np.ndarray:
         return (self._ub_masked + self._lb_masked) / 2 + (
             self._ub_masked - self._lb_masked
         ) / 2 * np.sin(x_optgra_masked)
 
-    def _eval_inv(self, x_masked: List[float]) -> List[float]:
+    def _eval_inv(self, x_masked: np.ndarray) -> np.ndarray:
         arg = (2 * x_masked - self._ub_masked - self._lb_masked) / (
             self._ub_masked - self._lb_masked
         )
-        # TODO: check that arg is in the interval [-1, 1]
-        return np.asin(arg)
+        if np.any((arg < -1.0) | (arg > 1.0)):
+            print(
+                "WARNING: Numerical inaccuracies encountered during KhanFunction inverse.",
+                "Clipping parameters to valid range.",
+            )
+            arg = np.clip(arg, -1.0, 1.0)
+        return np.arcsin(arg)
 
-    def _eval_grad(self, x_optgra_masked: List[float]) -> List[float]:
+    def _eval_grad(self, x_optgra_masked: np.ndarray) -> np.ndarray:
         return (self._ub_masked - self._lb_masked) / 2 * np.cos(x_optgra_masked)
 
-    def _eval_grad_inv(self, x_optgra_masked: List[float]) -> List[float]:
+    def _eval_grad_inv(self, x_optgra_masked: np.ndarray) -> np.ndarray:
         return -1 / (
             (self._lb_masked - self._ub_masked)
             * np.sqrt(
@@ -95,72 +99,74 @@ def _eval_grad_inv(self, x_optgra_masked: List[float]) -> List[float]:
             )
         )
 
-    def _apply_to_subset(self, x: List[float], func: Callable):
+    def _apply_to_subset(self, x: np.ndarray, func: Callable) -> np.ndarray:
         """Apply a given function only to a subset of x defined by self.mask."""
-        result = x.copy()  # Create a copy to preserve the original structure
-        result[self.mask] = func(x[self.mask])  # Apply the function only to the selected subset
+        # Create a copy to preserve the original structure
+        result = x.copy()
+        # Apply the function only to the selected subset
+        result[self.mask] = func(result[self.mask])
         return result
 
-    def eval(self, x_optgra: List[float]) -> List[float]:
+    def eval(self, x_optgra: np.ndarray) -> np.ndarray:
         """Convert :math:`x_{optgra}` to :math:`x`.
 
         Parameters
         ----------
-        x_optgra : List[float]
+        x_optgra : np.ndarray
             Decision vector passed to OPTGRA
 
         Returns
         -------
-        List[float]
+        np.ndarray
             Pagmo decision vector
         """
-        return self._apply_to_subset(x_optgra, self._eval)
+        return self._apply_to_subset(np.asarray(x_optgra), self._eval)
 
-    def eval_inv(self, x: List[float]) -> List[float]:
+    def eval_inv(self, x: np.ndarray) -> np.ndarray:
         """Convert :math:`x` to :math:`x_{optgra}`.
 
         Parameters
         ----------
-        x : List[float]
+        x : np.ndarray
             Pagmo decision vector
 
         Returns
         -------
-        List[float]
+        np.ndarray
             Decision vector passed to OPTGRA
 
         """
-        return self._apply_to_subset(x, self.eval_inv)
+        return self._apply_to_subset(np.asarray(x), self._eval_inv)
 
-    def eval_grad(self, x_optgra: List[float]) -> List[float]:
+    def eval_grad(self, x_optgra: np.ndarray) -> np.ndarray:
         """Gradient of ``eval`` function.
 
         Parameters
         ----------
-        x_optgra : List[float]
+        x_optgra : np.ndarray
             Decision vector passed to OPTGRA
 
         Returns
         -------
-        List[float]
+        np.ndarray
             Pagmo decision vector
         """
-        return self._apply_to_subset(x_optgra, self._eval_grad)
+        return np.diag(self._apply_to_subset(np.asarray(x_optgra), self._eval_grad))
 
-    def eval_grad_inv(self, x_optgra: List[float]) -> List[float]:
+    def eval_grad_inv(self, x_optgra: np.ndarray) -> np.ndarray:
         """Gradient of ``eval_inv`` method.
 
         Parameters
         ----------
-        x_optgra : List[float]
+        x_optgra : np.ndarray
             Decision vector passed to OPTGRA
 
         Returns
         -------
-        List[float]
+        np.ndarray
             Pagmo decision vector
         """
-        return self._apply_to_subset(x_optgra, self._eval_grad_inv)
+        return np.diag(self._apply_to_subset(np.asarray(x_optgra), self._eval_grad_inv))
 
 
 class optgra:
@@ -204,38 +210,38 @@ def _wrap_fitness_func(
         force_bounds: bool = False,
         khan_function: KhanFunction = None,
     ):
+        # get problem parameters
+        lb, ub = problem.get_bounds()
+
         def wrapped_fitness(x):
 
-            # if Khan function is used, we first need to convert to pagmo parameters
             if khan_function:
-                x = KhanFunction.eval(x)
-
-            fixed_x = x
-            lb, ub = problem.get_bounds()
+                # if Khan function is used, we first need to convert to pagmo parameters
+                x = khan_function.eval(x_optgra=x)
+            else:
+                # we are using vectorisation internally -> convert to ndarray
+                x = np.asarray(x)
 
             if force_bounds:
-                x = np.clip(x, lb, ub).tolist()
-                # for i in range(problem.get_nx()):
-                #     if x[i] < lb[i]:
-                #         fixed_x[i] = lb[i]
-                #     if x[i] > ub[i]:
-                #         fixed_x[i] = ub[i]
+                fixed_x = np.clip(x, lb, ub)
+            else:
+                fixed_x = x
 
-            result = deque(problem.fitness(fixed_x))
+            # call pagmo fitness function
+            result = problem.fitness(fixed_x)
 
             # add constraints derived from box bounds
             if bounds_to_constraints:
-                for i in range(len(lb)):
-                    if isfinite(lb[i]):
-                        result.append(x[i] - lb[i])
+                # Add (x[i] - lb[i]) for finite lb[i] and (ub[i] - x[i]) for finite ub[i]
+                result = np.concatenate(
+                    [result, (x - lb)[np.isfinite(lb)], (ub - x)[np.isfinite(ub)]]
+                )
 
-                for i in range(len(ub)):
-                    if isfinite(ub[i]):
-                        result.append(ub[i] - x[i])
+            # reorder constraint order, optgra expects the merit function last, pagmo has it first
+            # equivalent to rotating in a dequeue
+            result = np.concatenate([result[1:], result[0:1]])
 
-            # optgra expects the fitness last, pagmo has the fitness first
-            result.rotate(-1)
-            return list(result)
+            return result.tolist()  # return a list
 
         return wrapped_fitness
 
@@ -246,61 +252,64 @@ def _wrap_gradient_func(
         force_bounds=False,
         khan_function: KhanFunction = None,
     ):
-
+        # get the sparsity pattern to index the sparse gradients
         sparsity_pattern = problem.gradient_sparsity()
+        f_indices, x_indices = sparsity_pattern.T  # Unpack indices
 
+        # expected shape of the non-sparse gradient matrix
         shape = (problem.get_nf(), problem.get_nx())
 
+        # get problem parameters
+        lb, ub = problem.get_bounds()
+        nx = problem.get_nx()
+
         def wrapped_gradient(x):
 
-            # if Khan function is used, we first need to convert to pagmo parameters
             if khan_function:
-                x = KhanFunction.eval(x)
-
-            lb, ub = problem.get_bounds()
-            fixed_x = x
+                # if Khan function is used, we first need to convert to pagmo parameters
+                fixed_x = khan_function.eval(x_optgra=x)
+            else:
+                # we are using vectorisation internally -> convert to ndarray
+                fixed_x = np.asarray(x)
 
+            # force parameters to lower and upper bounds if needed
             if force_bounds:
-                np.clip(x, lb, ub).tolist()
-                # for i in range(problem.get_nx()):
-                #     if x[i] < lb[i]:
-                #         fixed_x[i] = lb[i]
-                #     if x[i] > ub[i]:
-                #         fixed_x[i] = ub[i]
+                fixed_x = np.clip(fixed_x, lb, ub)
 
+            # call the problem gradient function to retrieve sparse values
+            # gives derivative of merit function and constraints w.r.t. pagmo parameters
             sparse_values = problem.gradient(fixed_x)
 
-            nnz = len(sparse_values)
-
-            result = [[0 for j in range(shape[1])] for i in range(shape[0])]
+            # initialize non-sparse gradient matrix
+            result = np.zeros(shape)
 
             # expand gradient to dense representation
-            for i in range(nnz):
-                fIndex, xIndex = sparsity_pattern[i]
-
-                result[fIndex][xIndex] = sparse_values[i]
+            result[f_indices, x_indices] = sparse_values
 
             # add box-derived constraints
             if bounds_to_constraints:
-                lb, ub = problem.get_bounds()
-                for i in range(problem.get_nx()):
-                    if isfinite(lb[i]):
-                        box_bound_grad = [0 for j in range(problem.get_nx())]
-                        box_bound_grad[i] = 1
-                        result.append(box_bound_grad)
-
-                for i in range(len(ub)):
-                    if isfinite(ub[i]):
-                        box_bound_grad = [0 for j in range(problem.get_nx())]
-                        box_bound_grad[i] = -1
-                        result.append(box_bound_grad)
+
+                # lower bound gradients
+                finite_indices = np.isfinite(lb)  # Boolean mask for valid indices
+                box_lb_grads = np.eye(nx)[finite_indices]
+
+                # upper bound gradients
+                finite_indices = np.isfinite(ub)  # Boolean mask for valid indices
+                box_ub_grads = np.eye(nx)[finite_indices]
+
+                # append box bounds to gradient matrix
+                result = np.concatenate([result, box_lb_grads, box_ub_grads])
 
             # reorder constraint order, optgra expects the merit function last, pagmo has it first
             # equivalent to rotating in a dequeue
-            gradient_of_merit = result.pop(0)
-            result.append(gradient_of_merit)
+            result = np.vstack([result[1:], result[0]])
+
+            # if Khan function is used, we need to post multiply with the Khan function gradients
+            if khan_function:
+                khan_grad = khan_function.eval_grad(x)
+                result = result @ khan_grad
 
-            return result
+            return result.tolist()  # return as a list, not ndarray
 
         return wrapped_gradient
 
@@ -631,7 +640,7 @@ def evolve(self, population):
 
         if self.force_bounds:
             lb, ub = problem.get_bounds()
-            best_x = np.clip(best_x, lb, ub).tolist()
+            best_x = np.clip(best_x, lb, ub)
             # for i in range(problem.get_nx()):
             #     if best_x[i] < lb[i]:
             #         best_x[i] = lb[i]
diff --git a/tests/python/test.py b/tests/python/test.py
index 3148446..5c8e31f 100644
--- a/tests/python/test.py
+++ b/tests/python/test.py
@@ -144,26 +144,27 @@ def get_nic(self):
 
 class optgra_test(unittest.TestCase):
     def runTest(self):
-        self.constructor_test()
-        self.evolve_input_check_test()
-        self.basic_no_gradient_test()
-        self.gradient_no_constraints_test()
-        self.gradient_with_constraints_test()
-        self.box_constraints_test()
-        self.archipelago_evolve_test()
-        self.archipelago_pickle_test()
-        self.prepare_sensitivity_input_check_test()
-        self.prepare_sensitivity_test()
-        self.sensitivity_matrices_test()
-        self.sensitivity_new_callable_test()
-        self.sensitivity_constraint_delta_test()
-        self.sensitivity_active_constraints_test()
-        self.force_bounds_test()
-        self.force_bounds_fitness_test()
-        self.force_bounds_gradient_test()
-        self.get_name_test()
-        self.get_extra_info_test()
-        self.verbosity_test()
+        # self.constructor_test()
+        # self.evolve_input_check_test()
+        # self.basic_no_gradient_test()
+        # self.gradient_no_constraints_test()
+        # self.gradient_with_constraints_test()
+        # self.box_constraints_test()
+        # self.archipelago_evolve_test()
+        # self.archipelago_pickle_test()
+        # self.prepare_sensitivity_input_check_test()
+        # self.prepare_sensitivity_test()
+        # self.sensitivity_matrices_test()
+        # self.sensitivity_new_callable_test()
+        # self.sensitivity_constraint_delta_test()
+        # self.sensitivity_active_constraints_test()
+        # self.force_bounds_test()
+        self.khan_bounds_test()
+        # self.force_bounds_fitness_test()
+        # self.force_bounds_gradient_test()
+        # self.get_name_test()
+        # self.get_extra_info_test()
+        # self.verbosity_test()
 
     def constructor_test(self):
         # Check that invalid optimization method is rejected
@@ -676,7 +677,7 @@ def get_nic(self):
                 return 1
 
         # check bounds violation with normal optgra
-        algo = pygmo.algorithm(pyoptgra.optgra(force_bounds=False))
+        algo = pygmo.algorithm(pyoptgra.optgra(khan_bounds=False))
         prob = pygmo.problem(_prob_bound_test_no_gradient())
         pop = pygmo.population(prob, size=1)
         extracted = pop.problem.extract(_prob_bound_test_no_gradient)
@@ -685,8 +686,8 @@ def get_nic(self):
         extracted = pop.problem.extract(_prob_bound_test_no_gradient)
         self.assertTrue(extracted._bounds_violated)
 
-        # check bounds are forced when setting the argument
-        algo = pygmo.algorithm(pyoptgra.optgra(force_bounds=True))
+        # check bounds are satisfied when setting the argument
+        algo = pygmo.algorithm(pyoptgra.optgra(khan_bounds=True))
         prob = pygmo.problem(_prob_bound_test())
         pop = pygmo.population(prob, size=1)
         extracted = pop.problem.extract(_prob_bound_test)
@@ -696,7 +697,7 @@ def get_nic(self):
         self.assertFalse(extracted._bounds_violated)
 
         # check that population has valid members
-        algo = pygmo.algorithm(pyoptgra.optgra(force_bounds=True, bounds_to_constraints=False))
+        algo = pygmo.algorithm(pyoptgra.optgra(khan_bounds=True, bounds_to_constraints=False))
         prob = pygmo.problem(_prob_bound_test_no_gradient())
         pop = pygmo.population(prob, size=0)
         pop.push_back([2.47192039, -1.45880516, -9.03600606, -9.33306356, 3.85509973])

From 1dcbc62911728a65b5f97acdac5cb0ab4d75abc8 Mon Sep 17 00:00:00 2001
From: wmartens <waldemar.martens@esa.int>
Date: Fri, 31 Jan 2025 14:15:56 +0000
Subject: [PATCH 04/10] added unity gradient

---
 pyoptgra/__init__.py |   2 +-
 pyoptgra/_about.py   |   2 +-
 pyoptgra/optgra.py   | 142 ++++++++++++++++++++++++++-----------------
 tests/python/test.py |  37 ++++++++++-
 4 files changed, 125 insertions(+), 58 deletions(-)

diff --git a/pyoptgra/__init__.py b/pyoptgra/__init__.py
index e3886db..365812d 100644
--- a/pyoptgra/__init__.py
+++ b/pyoptgra/__init__.py
@@ -12,7 +12,7 @@
 # file, you can obtain them at https://www.gnu.org/licenses/gpl-3.0.txt
 # and https://essr.esa.int/license/european-space-agency-community-license-v2-4-weak-copyleft
 
-from .optgra import optgra  # noqa
+from .optgra import optgra, khan_function  # noqa
 
 try:
     from ._version import version as __version__  # noqa
diff --git a/pyoptgra/_about.py b/pyoptgra/_about.py
index 1ae62de..66336f1 100644
--- a/pyoptgra/_about.py
+++ b/pyoptgra/_about.py
@@ -1,3 +1,3 @@
 # Copyright (C) 2021 ESA, Flight Dynamics OPS-GF
 """Version number of pyoptgra package"""
-__version__ = "0.1.9"
+__version__ = "1.0.0"
diff --git a/pyoptgra/optgra.py b/pyoptgra/optgra.py
index 7d1cf8b..49e4cbc 100644
--- a/pyoptgra/optgra.py
+++ b/pyoptgra/optgra.py
@@ -28,19 +28,19 @@
 )
 
 
-class KhanFunction:
+class khan_function:
     """Function to smothly enforce optimisation parameter bounds as Michal Khan used to do:
 
     .. math::
 
-        x = \frac{x_{max} + x_{min}}{2} + \frac{x_{max} - x_{min}}{2} \cdot \sin(x_{optgra})
+        x = \frac{x_{max} + x_{min}}{2} + \frac{x_{max} - x_{min}}{2} \cdot \sin(x_{khan})
 
-    Where :math:`x` is the pagmo decision vector and :math:`x_{optgra}` is the decision vector
+    Where :math:`x` is the pagmo decision vector and :math:`x_{khan}` is the decision vector
     passed to OPTGRA. In this way parameter bounds are guaranteed to be satisfied, but the gradients
     near the bounds approach zero.
     """
 
-    def __init__(self, lb: List[float], ub: List[float]):
+    def __init__(self, lb: List[float], ub: List[float], unity_gradient: bool = True):
         """Constructor
 
         Parameters
@@ -49,9 +49,15 @@ def __init__(self, lb: List[float], ub: List[float]):
             Lower pagmo parameter bounds
         ub : List[float]
             Upper pagmo parameter bounds
+        unity_gradient : bool, optional
+            Uses an internal scaling that ensures that the derivative of pagmo parameters w.r.t.
+            khan parameters are unity at (lb + ub)/2. By default True.
+            Otherwise, the original Khan method is used that can result in strongly modified
+            gradients
         """
         self._lb = np.asarray(lb)
         self._ub = np.asarray(ub)
+        self._nx = len(lb)
 
         # determine finite lower/upper bounds
         finite_lb = np.isfinite(self._lb)
@@ -67,13 +73,21 @@ def __init__(self, lb: List[float], ub: List[float]):
             )
         # store the mask of finite bounds
         self.mask = finite_ub
-        self._lb_masked = lb[self.mask]
-        self._ub_masked = ub[self.mask]
+        self._lb_masked = self._lb[self.mask]
+        self._ub_masked = self._ub[self.mask]
+
+        # determine coefficients inside the sin function
+        self._a = 2 / (self._ub_masked - self._lb_masked) if unity_gradient else 1.0
+        self._b = (
+            -(self._ub_masked + self._lb_masked) / (self._ub_masked - self._lb_masked)
+            if unity_gradient
+            else 0.0
+        )
 
-    def _eval(self, x_optgra_masked: np.ndarray) -> np.ndarray:
+    def _eval(self, x_khan_masked: np.ndarray) -> np.ndarray:
         return (self._ub_masked + self._lb_masked) / 2 + (
             self._ub_masked - self._lb_masked
-        ) / 2 * np.sin(x_optgra_masked)
+        ) / 2 * np.sin(x_khan_masked * self._a + self._b)
 
     def _eval_inv(self, x_masked: np.ndarray) -> np.ndarray:
         arg = (2 * x_masked - self._ub_masked - self._lb_masked) / (
@@ -81,38 +95,49 @@ def _eval_inv(self, x_masked: np.ndarray) -> np.ndarray:
         )
         if np.any((arg < -1.0) | (arg > 1.0)):
             print(
-                "WARNING: Numerical inaccuracies encountered during KhanFunction inverse.",
+                "WARNING: Numerical inaccuracies encountered during khan_function inverse.",
                 "Clipping parameters to valid range.",
             )
             arg = np.clip(arg, -1.0, 1.0)
-        return np.arcsin(arg)
+        return (np.arcsin(arg) - self._b) / self._a
 
-    def _eval_grad(self, x_optgra_masked: np.ndarray) -> np.ndarray:
-        return (self._ub_masked - self._lb_masked) / 2 * np.cos(x_optgra_masked)
+    def _eval_grad(self, x_khan_masked: np.ndarray) -> np.ndarray:
+        return (
+            (self._ub_masked - self._lb_masked)
+            / 2
+            * np.cos(self._a * x_khan_masked + self._b)
+            * self._a
+        )
 
-    def _eval_grad_inv(self, x_optgra_masked: np.ndarray) -> np.ndarray:
-        return -1 / (
-            (self._lb_masked - self._ub_masked)
-            * np.sqrt(
-                ((self._lb_masked - x_optgra_masked) * (x_optgra_masked - self._ub_masked))
-                / (self._ub_masked - self._lb_masked) ** 2
+    def _eval_inv_grad(self, x_masked: np.ndarray) -> np.ndarray:
+        return (
+            -1
+            / self._a
+            / (
+                (self._lb_masked - self._ub_masked)
+                * np.sqrt(
+                    ((self._lb_masked - x_masked) * (x_masked - self._ub_masked))
+                    / (self._ub_masked - self._lb_masked) ** 2
+                )
             )
         )
 
-    def _apply_to_subset(self, x: np.ndarray, func: Callable) -> np.ndarray:
+    def _apply_to_subset(
+        self, x: np.ndarray, func: Callable, default_result: np.ndarray = None
+    ) -> np.ndarray:
         """Apply a given function only to a subset of x defined by self.mask."""
         # Create a copy to preserve the original structure
-        result = x.copy()
+        result = default_result if default_result is not None else x.copy()
         # Apply the function only to the selected subset
-        result[self.mask] = func(result[self.mask])
+        result[self.mask] = func(x[self.mask])
         return result
 
-    def eval(self, x_optgra: np.ndarray) -> np.ndarray:
+    def eval(self, x_khan: np.ndarray) -> np.ndarray:
         """Convert :math:`x_{optgra}` to :math:`x`.
 
         Parameters
         ----------
-        x_optgra : np.ndarray
+        x_khan : np.ndarray
             Decision vector passed to OPTGRA
 
         Returns
@@ -120,7 +145,7 @@ def eval(self, x_optgra: np.ndarray) -> np.ndarray:
         np.ndarray
             Pagmo decision vector
         """
-        return self._apply_to_subset(np.asarray(x_optgra), self._eval)
+        return self._apply_to_subset(np.asarray(x_khan), self._eval)
 
     def eval_inv(self, x: np.ndarray) -> np.ndarray:
         """Convert :math:`x` to :math:`x_{optgra}`.
@@ -138,12 +163,12 @@ def eval_inv(self, x: np.ndarray) -> np.ndarray:
         """
         return self._apply_to_subset(np.asarray(x), self._eval_inv)
 
-    def eval_grad(self, x_optgra: np.ndarray) -> np.ndarray:
+    def eval_grad(self, x_khan: np.ndarray) -> np.ndarray:
         """Gradient of ``eval`` function.
 
         Parameters
         ----------
-        x_optgra : np.ndarray
+        x_khan : np.ndarray
             Decision vector passed to OPTGRA
 
         Returns
@@ -151,22 +176,24 @@ def eval_grad(self, x_optgra: np.ndarray) -> np.ndarray:
         np.ndarray
             Pagmo decision vector
         """
-        return np.diag(self._apply_to_subset(np.asarray(x_optgra), self._eval_grad))
+        return np.diag(
+            self._apply_to_subset(np.asarray(x_khan), self._eval_grad, np.ones(self._nx))
+        )
 
-    def eval_grad_inv(self, x_optgra: np.ndarray) -> np.ndarray:
+    def eval_inv_grad(self, x: np.ndarray) -> np.ndarray:
         """Gradient of ``eval_inv`` method.
 
         Parameters
         ----------
-        x_optgra : np.ndarray
-            Decision vector passed to OPTGRA
+        x : np.ndarray
+            Pagmo decision vector
 
         Returns
         -------
         np.ndarray
-            Pagmo decision vector
+            Decision vector passed to OPTGRA
         """
-        return np.diag(self._apply_to_subset(np.asarray(x_optgra), self._eval_grad_inv))
+        return np.diag(self._apply_to_subset(np.asarray(x), self._eval_inv_grad, np.ones(self._nx)))
 
 
 class optgra:
@@ -208,19 +235,19 @@ def _wrap_fitness_func(
         problem,
         bounds_to_constraints: bool = True,
         force_bounds: bool = False,
-        khan_function: KhanFunction = None,
+        khanf: khan_function = None,
     ):
         # get problem parameters
         lb, ub = problem.get_bounds()
 
         def wrapped_fitness(x):
 
-            if khan_function:
+            # we are using vectorisation internally -> convert to ndarray
+            x = np.asarray(x, dtype=np.float64)
+
+            if khanf:
                 # if Khan function is used, we first need to convert to pagmo parameters
-                x = khan_function.eval(x_optgra=x)
-            else:
-                # we are using vectorisation internally -> convert to ndarray
-                x = np.asarray(x)
+                x = khanf.eval(x_khan=x)
 
             if force_bounds:
                 fixed_x = np.clip(x, lb, ub)
@@ -250,7 +277,7 @@ def _wrap_gradient_func(
         problem,
         bounds_to_constraints: bool = True,
         force_bounds=False,
-        khan_function: KhanFunction = None,
+        khanf: khan_function = None,
     ):
         # get the sparsity pattern to index the sparse gradients
         sparsity_pattern = problem.gradient_sparsity()
@@ -265,16 +292,18 @@ def _wrap_gradient_func(
 
         def wrapped_gradient(x):
 
-            if khan_function:
+            # we are using vectorisation internally -> convert to ndarray
+            x = np.asarray(x, dtype=np.float64)
+
+            if khanf:
                 # if Khan function is used, we first need to convert to pagmo parameters
-                fixed_x = khan_function.eval(x_optgra=x)
-            else:
-                # we are using vectorisation internally -> convert to ndarray
-                fixed_x = np.asarray(x)
+                x = khanf.eval(x_khan=x)
 
             # force parameters to lower and upper bounds if needed
             if force_bounds:
-                fixed_x = np.clip(fixed_x, lb, ub)
+                fixed_x = np.clip(x, lb, ub)
+            else:
+                fixed_x = x
 
             # call the problem gradient function to retrieve sparse values
             # gives derivative of merit function and constraints w.r.t. pagmo parameters
@@ -305,8 +334,8 @@ def wrapped_gradient(x):
             result = np.vstack([result[1:], result[0]])
 
             # if Khan function is used, we need to post multiply with the Khan function gradients
-            if khan_function:
-                khan_grad = khan_function.eval_grad(x)
+            if khanf:
+                khan_grad = khanf.eval_grad(x)
                 result = result @ khan_grad
 
             return result.tolist()  # return as a list, not ndarray
@@ -378,11 +407,14 @@ def __init__(
 
                 .. math::
 
-                    x = \frac{x_{max} + x_{min}}{2} + \frac{x_{max} - x_{min}}{2} \cdot \sin(x_{optgra})
+                    x = \frac{x_{max} + x_{min}}{2} + \frac{x_{max} - x_{min}}{2} \cdot \sin(x_{Khan})
 
-                Where :math:`x` is the pagmo decision vector and :math:`x_{optgra}` is the decision
+                Where :math:`x` is the pagmo decision vector and :math:`x_{Khan}` is the decision
                 vector passed to OPTGRA. In this way parameter bounds are guaranteed to be
                 satisfied, but the gradients near the bounds approach zero. By default False.
+                Pyoptgra uses a variant of the above method that additionally scales the
+                argument of the :math:`\sin` function such that the derivatives
+                :math:`\frac{x_{Khan}}{x}` are unity in the center of the box bounds.
             optimization_method: select 0 for steepest descent, 1 for modified spectral conjugate
                 gradient method, 2 for spectral conjugate gradient method and 3 for conjugate
                 gradient method
@@ -564,16 +596,16 @@ def evolve(self, population):
         idx = list(population.get_ID()).index(selected[0][0])
 
         # optional Khan function to enforce parameter bounds
-        khan_function = KhanFunction(*problem.get_bounds()) if self.khan_bounds else None
+        khanf = khan_function(*problem.get_bounds()) if self.khan_bounds else None
 
         fitness_func = optgra._wrap_fitness_func(
-            problem, self.bounds_to_constraints, self.force_bounds, khan_function
+            problem, self.bounds_to_constraints, self.force_bounds, khanf
         )
         grad_func = None
         derivatives_computation = 2
         if problem.has_gradient():
             grad_func = optgra._wrap_gradient_func(
-                problem, self.bounds_to_constraints, self.force_bounds, khan_function
+                problem, self.bounds_to_constraints, self.force_bounds, khanf
             )
             derivatives_computation = 1
 
@@ -610,7 +642,7 @@ def evolve(self, population):
         # get initial x
         x0 = population.get_x()[idx]
         result = optimize(
-            initial_x=khan_function.eval_inv(x0) if khan_function else x0,
+            initial_x=khanf.eval_inv(x0) if khanf else x0,
             constraint_types=constraint_types,
             fitness_callback=fitness_func,
             gradient_callback=grad_func,
@@ -635,8 +667,8 @@ def evolve(self, population):
         best_x, best_f, finopt = result
 
         # if a Khan function is used we first need to convert to pagmo parameters
-        if khan_function:
-            best_x = khan_function.eval(best_x)
+        if khanf:
+            best_x = khanf.eval(best_x)
 
         if self.force_bounds:
             lb, ub = problem.get_bounds()
diff --git a/tests/python/test.py b/tests/python/test.py
index 5c8e31f..dcbecec 100644
--- a/tests/python/test.py
+++ b/tests/python/test.py
@@ -14,6 +14,8 @@
 
 import unittest
 
+import numpy as np
+
 import pygmo
 
 import pyoptgra
@@ -159,7 +161,8 @@ def runTest(self):
         # self.sensitivity_constraint_delta_test()
         # self.sensitivity_active_constraints_test()
         # self.force_bounds_test()
-        self.khan_bounds_test()
+        # self.khan_bounds_test()
+        self.khan_function_test()
         # self.force_bounds_fitness_test()
         # self.force_bounds_gradient_test()
         # self.get_name_test()
@@ -711,6 +714,38 @@ def get_nic(self):
             self.assertTrue(pop.champion_x[i] >= lb[i])
             self.assertTrue(pop.champion_x[i] <= ub[i])
 
+    def khan_function_test(self):
+
+        for unity_gradient in [True, False]:  # test both variants
+            lb = [-10, 0, -np.inf, -np.inf, -20]
+            ub = [10, 30, np.inf, -np.inf, -10]
+            kfun = pyoptgra.khan_function(lb, ub, unity_gradient)
+
+            # check function and its inversion
+            x = np.asarray([-1, 2, 4, 4, -15.0])
+            x_optgra = kfun.eval_inv(x)
+            x_check = kfun.eval(x_optgra)
+            self.assertListEqual(x.tolist(), x_check.tolist())
+
+            # check gradient and its inversion
+            dx_dxog = kfun.eval_inv_grad(x)
+            dxog_dx = kfun.eval_grad(x_optgra)
+            check_mat = dx_dxog @ dxog_dx  # expect unity matrix
+            np.testing.assert_allclose(check_mat, np.eye(5), atol=1e-10)
+
+            # compare with numerical gradient
+            dx_dxog_num = pygmo.estimate_gradient_h(lambda _x: kfun.eval_inv(_x), x).reshape(5, 5)
+            dxog_dx_num = pygmo.estimate_gradient_h(lambda _x: kfun.eval(_x), x_optgra).reshape(
+                5, 5
+            )
+            np.testing.assert_allclose(dx_dxog_num, dx_dxog, atol=1e-7)
+            np.testing.assert_allclose(dxog_dx_num, dxog_dx, atol=1e-7)
+
+            # one-sided bound is not supported
+            ub = [10, 30, np.inf, -np.inf, np.inf]
+            with self.assertRaises(ValueError):
+                pyoptgra.khan_function(lb, ub)
+
     def get_name_test(self):
         algo = pygmo.algorithm(pyoptgra.optgra())
         self.assertEqual(algo.get_name(), "Optgra")

From 6104295453db62b4683ac966b7d0e1c0a277491c Mon Sep 17 00:00:00 2001
From: wmartens <waldemar.martens@esa.int>
Date: Fri, 31 Jan 2025 14:17:10 +0000
Subject: [PATCH 05/10] activated tests again

---
 tests/python/test.py | 42 +++++++++++++++++++++---------------------
 1 file changed, 21 insertions(+), 21 deletions(-)

diff --git a/tests/python/test.py b/tests/python/test.py
index dcbecec..ddb6cd8 100644
--- a/tests/python/test.py
+++ b/tests/python/test.py
@@ -146,28 +146,28 @@ def get_nic(self):
 
 class optgra_test(unittest.TestCase):
     def runTest(self):
-        # self.constructor_test()
-        # self.evolve_input_check_test()
-        # self.basic_no_gradient_test()
-        # self.gradient_no_constraints_test()
-        # self.gradient_with_constraints_test()
-        # self.box_constraints_test()
-        # self.archipelago_evolve_test()
-        # self.archipelago_pickle_test()
-        # self.prepare_sensitivity_input_check_test()
-        # self.prepare_sensitivity_test()
-        # self.sensitivity_matrices_test()
-        # self.sensitivity_new_callable_test()
-        # self.sensitivity_constraint_delta_test()
-        # self.sensitivity_active_constraints_test()
-        # self.force_bounds_test()
-        # self.khan_bounds_test()
+        self.constructor_test()
+        self.evolve_input_check_test()
+        self.basic_no_gradient_test()
+        self.gradient_no_constraints_test()
+        self.gradient_with_constraints_test()
+        self.box_constraints_test()
+        self.archipelago_evolve_test()
+        self.archipelago_pickle_test()
+        self.prepare_sensitivity_input_check_test()
+        self.prepare_sensitivity_test()
+        self.sensitivity_matrices_test()
+        self.sensitivity_new_callable_test()
+        self.sensitivity_constraint_delta_test()
+        self.sensitivity_active_constraints_test()
+        self.force_bounds_test()
+        self.khan_bounds_test()
         self.khan_function_test()
-        # self.force_bounds_fitness_test()
-        # self.force_bounds_gradient_test()
-        # self.get_name_test()
-        # self.get_extra_info_test()
-        # self.verbosity_test()
+        self.force_bounds_fitness_test()
+        self.force_bounds_gradient_test()
+        self.get_name_test()
+        self.get_extra_info_test()
+        self.verbosity_test()
 
     def constructor_test(self):
         # Check that invalid optimization method is rejected

From f78961386a2db560453b54ea66bcdb3161437606 Mon Sep 17 00:00:00 2001
From: wmartens <waldemar.martens@esa.int>
Date: Fri, 31 Jan 2025 14:41:47 +0000
Subject: [PATCH 06/10] fixing pipeline

---
 pyoptgra/optgra.py   | 13 ++++++-------
 tests/python/test.py |  2 +-
 2 files changed, 7 insertions(+), 8 deletions(-)

diff --git a/pyoptgra/optgra.py b/pyoptgra/optgra.py
index 49e4cbc..f91ea4c 100644
--- a/pyoptgra/optgra.py
+++ b/pyoptgra/optgra.py
@@ -12,9 +12,8 @@
 # file, you can obtain them at https://www.gnu.org/licenses/gpl-3.0.txt
 # and https://essr.esa.int/license/european-space-agency-community-license-v2-4-weak-copyleft
 
-from collections import deque
 from math import isfinite
-from typing import List, Tuple, Union, Callable
+from typing import List, Tuple, Union, Callable, Optional
 import numpy as np
 
 from pygmo import s_policy, select_best
@@ -38,7 +37,7 @@ class khan_function:
     Where :math:`x` is the pagmo decision vector and :math:`x_{khan}` is the decision vector
     passed to OPTGRA. In this way parameter bounds are guaranteed to be satisfied, but the gradients
     near the bounds approach zero.
-    """
+    """  # noqa: W605
 
     def __init__(self, lb: List[float], ub: List[float], unity_gradient: bool = True):
         """Constructor
@@ -123,7 +122,7 @@ def _eval_inv_grad(self, x_masked: np.ndarray) -> np.ndarray:
         )
 
     def _apply_to_subset(
-        self, x: np.ndarray, func: Callable, default_result: np.ndarray = None
+        self, x: np.ndarray, func: Callable, default_result: Optional[np.ndarray] = None
     ) -> np.ndarray:
         """Apply a given function only to a subset of x defined by self.mask."""
         # Create a copy to preserve the original structure
@@ -235,7 +234,7 @@ def _wrap_fitness_func(
         problem,
         bounds_to_constraints: bool = True,
         force_bounds: bool = False,
-        khanf: khan_function = None,
+        khanf: Optional[khan_function] = None,
     ):
         # get problem parameters
         lb, ub = problem.get_bounds()
@@ -277,7 +276,7 @@ def _wrap_gradient_func(
         problem,
         bounds_to_constraints: bool = True,
         force_bounds=False,
-        khanf: khan_function = None,
+        khanf: Optional[khan_function] = None,
     ):
         # get the sparsity pattern to index the sparse gradients
         sparsity_pattern = problem.gradient_sparsity()
@@ -428,7 +427,7 @@ def __init__(
             max_iterations, max_correction_iterations, max_distance_per_iteration,
                 or perturbation_for_snd_order_derivatives are negative
 
-        """
+        """  # noqa: W605
         self.max_iterations = max_iterations
         self.max_correction_iterations = max_correction_iterations
         self.max_distance_per_iteration = max_distance_per_iteration
diff --git a/tests/python/test.py b/tests/python/test.py
index ddb6cd8..b036a27 100644
--- a/tests/python/test.py
+++ b/tests/python/test.py
@@ -725,7 +725,7 @@ def khan_function_test(self):
             x = np.asarray([-1, 2, 4, 4, -15.0])
             x_optgra = kfun.eval_inv(x)
             x_check = kfun.eval(x_optgra)
-            self.assertListEqual(x.tolist(), x_check.tolist())
+            np.testing.assert_allclose(x, x_check, atol=1e-10)
 
             # check gradient and its inversion
             dx_dxog = kfun.eval_inv_grad(x)

From d2ed9627d6375208f8db0e7f977efa45a1a84544 Mon Sep 17 00:00:00 2001
From: wmartens <waldemar.martens@esa.int>
Date: Fri, 31 Jan 2025 14:46:06 +0000
Subject: [PATCH 07/10] fixed static checks

---
 pyoptgra/__init__.py | 2 +-
 pyoptgra/optgra.py   | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/pyoptgra/__init__.py b/pyoptgra/__init__.py
index 365812d..52ef013 100644
--- a/pyoptgra/__init__.py
+++ b/pyoptgra/__init__.py
@@ -12,7 +12,7 @@
 # file, you can obtain them at https://www.gnu.org/licenses/gpl-3.0.txt
 # and https://essr.esa.int/license/european-space-agency-community-license-v2-4-weak-copyleft
 
-from .optgra import optgra, khan_function  # noqa
+from .optgra import khan_function, optgra  # noqa
 
 try:
     from ._version import version as __version__  # noqa
diff --git a/pyoptgra/optgra.py b/pyoptgra/optgra.py
index f91ea4c..247b88b 100644
--- a/pyoptgra/optgra.py
+++ b/pyoptgra/optgra.py
@@ -13,9 +13,9 @@
 # and https://essr.esa.int/license/european-space-agency-community-license-v2-4-weak-copyleft
 
 from math import isfinite
-from typing import List, Tuple, Union, Callable, Optional
-import numpy as np
+from typing import Callable, List, Optional, Tuple, Union
 
+import numpy as np
 from pygmo import s_policy, select_best
 
 from .core import (

From 2ce78d60c917ddb8936ff50acfde847bfd627a7b Mon Sep 17 00:00:00 2001
From: wmartens <waldemar.martens@esa.int>
Date: Fri, 31 Jan 2025 14:53:07 +0000
Subject: [PATCH 08/10] finished tests

---
 tests/python/test.py | 27 +++++++++++++++++++++++++++
 1 file changed, 27 insertions(+)

diff --git a/tests/python/test.py b/tests/python/test.py
index b036a27..915d944 100644
--- a/tests/python/test.py
+++ b/tests/python/test.py
@@ -306,6 +306,7 @@ def gradient_no_constraints_test(self):
         self.assertLess(new_best, previous_best)
 
     def gradient_with_constraints_test(self):
+        # 1. Run Luksan-Vlcek problem with optgra
         prob = pygmo.problem(luksan_vlcek())
         prob.c_tol = 1e-6
         og = pyoptgra.optgra(
@@ -333,6 +334,32 @@ def gradient_with_constraints_test(self):
         for i in [5, 6]:
             self.assertLess(pop.champion_f[i], 1e-6)
 
+        # 2. Run the same test with khan_bounds
+        og = pyoptgra.optgra(
+            optimization_method=1,
+            max_iterations=100,
+            max_correction_iterations=100,
+            max_distance_per_iteration=10,
+            khan_bounds=True,
+        )
+        algo = pygmo.algorithm(og)
+        pop2 = pygmo.population(prob, size=0, seed=1)  # empty population
+        pop2.push_back([0.5, 0.5, -0.5, 0.4, 0.3, 0.7])  # add initial guess
+
+        # Calling optgra
+        pop2 = algo.evolve(pop2)  # run the optimisation
+
+        # objective function
+        self.assertLess(pop2.champion_f[0], 2.26)
+
+        # equality constraints
+        for i in [1, 2, 3, 4]:
+            self.assertAlmostEqual(pop2.champion_f[i], 0.0, 6)
+
+        # inequality constraints
+        for i in [5, 6]:
+            self.assertLess(pop2.champion_f[i], 1e-6)
+
     def constraints_with_default_tolerances_test(self):
         prob = pygmo.problem(luksan_vlcek())
         pop = pygmo.population(prob, size=0, seed=1)  # empty population

From 6b1e42fc49e3faa4d5e3d4f628341f2b1750555c Mon Sep 17 00:00:00 2001
From: wmartens <waldemar.martens@esa.int>
Date: Fri, 31 Jan 2025 15:03:37 +0000
Subject: [PATCH 09/10] testsing doctrings

---
 pyoptgra/optgra.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/pyoptgra/optgra.py b/pyoptgra/optgra.py
index 247b88b..5b1a0ae 100644
--- a/pyoptgra/optgra.py
+++ b/pyoptgra/optgra.py
@@ -359,7 +359,7 @@ def __init__(
         optimization_method: int = 2,
         log_level: int = 0,
     ) -> None:
-        """
+        r"""
         Initialize a wrapper instance for the OPTGRA algorithm.
 
         Some of the construction arguments, for example the scaling factors, depend on the dimension
@@ -413,7 +413,7 @@ def __init__(
                 satisfied, but the gradients near the bounds approach zero. By default False.
                 Pyoptgra uses a variant of the above method that additionally scales the
                 argument of the :math:`\sin` function such that the derivatives
-                :math:`\frac{x_{Khan}}{x}` are unity in the center of the box bounds.
+                :math:`\frac{d x_{Khan}}{d x}` are unity in the center of the box bounds.
             optimization_method: select 0 for steepest descent, 1 for modified spectral conjugate
                 gradient method, 2 for spectral conjugate gradient method and 3 for conjugate
                 gradient method

From d97522c24b127cd1d91ff75e800c6e1bd38b9314 Mon Sep 17 00:00:00 2001
From: wmartens <waldemar.martens@esa.int>
Date: Fri, 31 Jan 2025 15:19:43 +0000
Subject: [PATCH 10/10] bumped version to 1.0.0

---
 README.rst     | 2 +-
 pyproject.toml | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/README.rst b/README.rst
index 228a560..e17d62b 100644
--- a/README.rst
+++ b/README.rst
@@ -9,7 +9,7 @@ Optgra
 This repository provides *pyoptgra*, a python package wrapping (and including) OPTGRA.
 OPTGRA is an optimization algorithm developed and implemented by Johannes Schoenmaekers, it is specifically designed for near-linear constrained problems, which commonly occur in trajectory optimization.
 
-The full documentation can be found here_.
+The full documentation can be found here_ or on the `ESA-internal page <https://midas.io.esa.int/optgra-wrapper>`_.
 
 .. _here: https://esa.github.io/pyoptgra/
 
diff --git a/pyproject.toml b/pyproject.toml
index 4dc65b3..173cb58 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -12,7 +12,7 @@ license = {text = "GPL-3.0 or ESCL-2.4"}
 name = "pyoptgra"
 readme = "README.rst"
 requires-python = ">=3.9"
-version = "0.1.10"
+version = "1.0.0"
 
 [build-system]
 build-backend = "scikit_build_core.build"