Add functionality to Region class

src/estimagic/optimization/tranquilo/bounds.py

@@ -0,0 +1,27 @@
+from dataclasses import dataclass, replace
+
+import numpy as np
+
+
+@dataclass
+class Bounds:
+    """Parameter bounds."""
+
+    lower: np.ndarray
+    upper: np.ndarray
+
+    def __post_init__(self):
+        self.has_any = _any_finite(self.lower, self.upper)
+
+    # make it behave like a NamedTuple
+    def _replace(self, **kwargs):
+        return replace(self, **kwargs)
+
+
+def _any_finite(lb, ub):
+    out = False
+    if lb is not None and np.isfinite(lb).any():
+        out = True
+    if ub is not None and np.isfinite(ub).any():
+        out = True
+    return out

src/estimagic/optimization/tranquilo/estimate_variance.py

@@ -5,7 +5,7 @@
 
 from estimagic.optimization.tranquilo.get_component import get_component
 from estimagic.optimization.tranquilo.new_history import History
-from estimagic.optimization.tranquilo.options import Region
+from estimagic.optimization.tranquilo.region import Region
 
 
 def get_variance_estimator(fitter, user_options):

src/estimagic/optimization/tranquilo/geometry.py

@@ -2,12 +2,8 @@
 
 import numpy as np
 
-from estimagic.optimization.tranquilo.options import Region
-from estimagic.optimization.tranquilo.sample_points import (
-    _get_effective_bounds,
-    _map_from_feasible_trustregion,
-    get_sampler,
-)
+from estimagic.optimization.tranquilo.region import Region
+from estimagic.optimization.tranquilo.sample_points import get_sampler
 
 
 def get_geometry_checker_pair(
@@ -26,7 +22,7 @@ def get_geometry_checker_pair(
             samples drawn inside a box or a ball, respectively.
         n_params (int): Number of parameters.
         n_simulations (int): Number of simulations for the mean calculation.
-        bounds (Bounds): The parameter bounds.
+        bounds (Bounds): The parameter bounds. See module bounds.py.
 
     Returns:
         callable: The sample quality calculator.
@@ -45,7 +41,7 @@ def get_geometry_checker_pair(
 
     _checker = built_in_checker[checker]
 
-    quality_calculator = partial(_checker["quality_calculator"], bounds=bounds)
+    quality_calculator = _checker["quality_calculator"]
     cutoff_simulator = partial(
         _checker["cutoff_simulator"],
         reference_sampler=reference_sampler,
@@ -66,26 +62,26 @@ def log_d_cutoff_simulator(
         rng (np.random.Generator): The random number generator.
         reference_sampler (str): Either "box" or "ball", corresponding to comparison
             samples drawn inside a box or a ball, respectively.
-        bounds (Bounds): The parameter bounds.
+        bounds (Bounds): The parameter bounds. See module bounds.py.
         n_params (int): Dimensionality of the sample.
         n_simulations (int): Number of simulations for the mean calculation.
 
     Returns:
         float: The simulated mean logarithm of the d-optimality criterion.
 
     """
-    _sampler = get_sampler(reference_sampler, bounds)
-    trustregion = Region(center=np.zeros(n_params), radius=1, shape=None)
+    _sampler = get_sampler(reference_sampler)
+    trustregion = Region(center=np.zeros(n_params), radius=1, bounds=bounds)
     sampler = partial(_sampler, trustregion=trustregion)
     raw = []
     for _ in range(n_simulations):
         x = sampler(n_points=n_samples, rng=rng)
-        raw.append(log_d_quality_calculator(x, trustregion, bounds))
+        raw.append(log_d_quality_calculator(x, trustregion))
     out = np.nanmean(raw)
     return out
 
 
-def log_d_quality_calculator(sample, trustregion, bounds):
+def log_d_quality_calculator(sample, trustregion):
     """Logarithm of the d-optimality criterion.
 
     For a data sample x the log_d_criterion is defined as log(det(x.T @ x)). If the
@@ -94,15 +90,13 @@ def log_d_quality_calculator(sample, trustregion, bounds):
 
     Args:
         sample (np.ndarray): The data sample, shape = (n, p).
-        trustregion (TrustRegion): NamedTuple with attributes center and radius.
-        bounds (Bounds): The parameter bounds.
+        trustregion (Region): Trustregion. See module region.py.
 
     Returns:
         np.ndarray: The criterion values, shape = (n, ).
 
     """
-    effective_bounds = _get_effective_bounds(trustregion, bounds)
-    points = _map_from_feasible_trustregion(sample, effective_bounds)
+    points = trustregion.map_to_unit(sample)
     n_samples, n_params = points.shape
     xtx = points.T @ points
     det = np.linalg.det(xtx / n_samples)

src/estimagic/optimization/tranquilo/models.py

@@ -4,7 +4,7 @@
 import numpy as np
 from numba import njit
 
-from estimagic.optimization.tranquilo.options import Region
+from estimagic.optimization.tranquilo.region import Region
 
 
 @dataclass

src/estimagic/optimization/tranquilo/options.py

@@ -1,4 +1,3 @@
-from dataclasses import dataclass
 from typing import NamedTuple
 
 import numpy as np
@@ -9,17 +8,6 @@
 )
 
 
-@dataclass
-class Bounds:
-    """Stopping criteria."""
-
-    lower: np.ndarray
-    upper: np.ndarray
-
-    def __post_init__(self):
-        self.has_any = _check_if_there_are_bounds(self.lower, self.upper)
-
-
 class StopOptions(NamedTuple):
     """Criteria for stopping without successful convergence."""
 
@@ -54,12 +42,6 @@ class RadiusOptions(NamedTuple):
     max_radius_to_step_ratio: float = np.inf
 
 
-class Region(NamedTuple):
-    center: np.ndarray
-    radius: float
-    shape: str
-
-
 class AcceptanceOptions(NamedTuple):
     confidence_level: float = 0.8
     power_level: float = 0.8
@@ -76,12 +58,3 @@ class StagnationOptions(NamedTuple):
     sample_increment: int = 1
     max_trials: int = 1
     drop: bool = True
-
-
-def _check_if_there_are_bounds(lb, ub):
-    out = False
-    if lb is not None and np.isfinite(lb).any():
-        out = True
-    if ub is not None and np.isfinite(ub).any():
-        out = True
-    return out

src/estimagic/optimization/tranquilo/region.py

@@ -0,0 +1,116 @@
+from dataclasses import dataclass, replace
+
+import numpy as np
+
+from estimagic.optimization.tranquilo.bounds import Bounds
+from estimagic.optimization.tranquilo.volume import (
+    get_radius_of_cube_with_volume_of_sphere,
+)
+
+
+@dataclass
+class Region:
+    """Trust region."""
+
+    center: np.ndarray
+    radius: float
+    bounds: Bounds = None
+
+    @property
+    def shape(self) -> str:
+        any_bounds_binding = _any_bounds_binding(
+            bounds=self.bounds, center=self.center, radius=self.radius
+        )
+        return "cube" if any_bounds_binding else "sphere"
+
+    @property
+    def cube_bounds(self) -> Bounds:
+        if self.shape == "sphere":
+            raise AttributeError(
+                "The trustregion is a sphere, and thus has no cube bounds."
+            )
+        radius = get_radius_of_cube_with_volume_of_sphere(
+            self.radius, len(self.center), scaling_factor=1.0
+        )
+        bounds = _get_cube_bounds(center=self.center, radius=radius, bounds=self.bounds)
+        return bounds
+
+    @property
+    def cube_center(self) -> np.ndarray:
+        if self.shape == "sphere":
+            raise AttributeError(
+                "The trustregion is a sphere, and thus has no cube center."
+            )
+        cube_bounds = self.cube_bounds
+        center = (cube_bounds.lower + cube_bounds.upper) / 2
+        return center
+
+    def map_to_unit(self, x: np.ndarray) -> np.ndarray:
+        """Map points from the trustregion to the unit sphere or cube."""
+        if self.shape == "sphere":
+            out = _map_to_unit_sphere(x, center=self.center, radius=self.radius)
+        else:
+            out = _map_to_unit_cube(x, cube_bounds=self.cube_bounds)
+        return out
+
+    def map_from_unit(self, x: np.ndarray) -> np.ndarray:
+        """Map points from the unit sphere or cube to the trustregion."""
+        if self.shape == "sphere":
+            out = _map_from_unit_sphere(x, center=self.center, radius=self.radius)
+        else:
+            out = _map_from_unit_cube(x, cube_bounds=self.cube_bounds)
+        return out
+
+    # make it behave like a NamedTuple
+    def _replace(self, **kwargs):
+        return replace(self, **kwargs)
+
+
+def _map_to_unit_cube(x, cube_bounds):
+    """Map points from the trustregion to the unit cube."""
+    out = 2 * (x - cube_bounds.lower) / (cube_bounds.upper - cube_bounds.lower) - 1
+    return out
+
+
+def _map_to_unit_sphere(x, center, radius):
+    """Map points from the trustregion to the unit sphere."""
+    out = (x - center) / radius
+    return out
+
+
+def _map_from_unit_cube(x, cube_bounds):
+    """Map points from the unit cube to the trustregion."""
+    out = (cube_bounds.upper - cube_bounds.lower) * (x + 1) / 2 + cube_bounds.lower
+    return out
+
+
+def _map_from_unit_sphere(x, center, radius):
+    """Map points from the unit sphere to the trustregion."""
+    out = x * radius + center
+    return out
+
+
+def _get_cube_bounds(center, radius, bounds):
+    """Get new bounds that define the intersection of the trustregion and the bounds."""
+    lower_bounds = center - radius
+    upper_bounds = center + radius
+
+    if bounds is not None and bounds.lower is not None:
+        lower_bounds = np.clip(lower_bounds, bounds.lower, np.inf)
+
+    if bounds is not None and bounds.upper is not None:
+        upper_bounds = np.clip(upper_bounds, -np.inf, bounds.upper)
+
+    return Bounds(lower=lower_bounds, upper=upper_bounds)
+
+
+def _any_bounds_binding(bounds, center, radius):
+    """Check if any bound is binding, i.e. inside the trustregion."""
+    out = False
+    if bounds is not None and bounds.has_any:
+        if bounds.lower is not None:
+            lower_binding = np.min(center - bounds.lower) <= radius
+        if bounds.upper is not None:
+            upper_binding = np.min(bounds.upper - center) <= radius
+        out = lower_binding or upper_binding
+    return out