badr 0.1.0__py3-none-any.whl

badr/__base__.py

@@ -0,0 +1,126 @@
+from badr.algorithms import SLSQP
+from badr.datasets import Dataset
+from badr.metrics import FairnessMetric
+from badr.models import Model
+from badr.oracles import ImplicitOracle, StochasticOracle
+
+
+class Badr:
+    def __init__(
+        self,
+        dset: Dataset,
+        model: Model,
+        metric: FairnessMetric,
+        train_test: str = "train",
+        oracle: str = "implicit",
+        solver_cls=None,
+        solver_kwargs=None,
+    ) -> None:
+        """
+        Parameters
+        ----------
+        dset : Dataset
+            Dataset with (X_train, y_train), (X_test, y_test), and groups.
+        model : Model
+            Model with set_group_weights(...), fit(...), and coef_/intercept_.
+        metric : FairnessMetric
+            Metric; bound to `model` if `metric.model is None`.
+        train_test : {"train", "test"}, default="train"
+            Which split to use.
+        oracle : {"implicit", "stochastic"}, default="implicit"
+            Oracle implementation to use.
+        solver_cls : type, optional
+            Solver class (default: SLSQP).
+        solver_kwargs : dict, optional
+            Keyword args passed to `solver_cls(...)`.
+
+        Raises
+        ------
+        ValueError
+            If `oracle` is not one of {"implicit", "stochastic"}.
+        """
+        if metric.model is None:
+            metric.set_model(model)
+        self.dset = dset
+        self.model = model
+        self.metric = metric
+        self.train_test = train_test
+        self.X = dset.X_train if train_test == "train" else dset.X_test
+        self.y = dset.y_train if train_test == "train" else dset.y_test
+        if oracle == "implicit":
+            self.oracle = ImplicitOracle(
+                dset=dset,
+                model=model,
+                metric=metric,
+                train_test=train_test,
+            )
+        elif oracle == "stochastic":
+            self.oracle = StochasticOracle(
+                dset=dset,
+                model=model,
+                metric=metric,
+                train_test=train_test,
+            )
+        else:
+            raise ValueError(
+                f"Unknown oracle type: {oracle}. Supported types are 'implicit' and 'stochastic'."
+            )
+        self.solver_cls = solver_cls or SLSQP
+        self.solver_kwargs = solver_kwargs or {}
+        self._solver = None
+
+    def set_solver(self, solver):
+        """
+        Set a solver instance to use in `run`.
+
+        Parameters
+        ----------
+        solver
+            Solver instance. If `solver.oracle is None`, `run` will set it.
+
+        Returns
+        -------
+        Badr
+            Self.
+        """
+        self._solver = solver
+        return self
+
+    def run(self, **run_kwargs) -> None:
+        """
+        Run the solver, set group weights, refit the model, and compute outputs.
+
+        Parameters
+        ----------
+        **run_kwargs
+            Passed to `solver.run(**run_kwargs)`.
+
+        Sets Attributes
+        ---------------
+        group_weights
+            Learned group weights.
+        coef_
+            Fitted coefficients.
+        intercept_
+            Fitted intercept.
+        fairness
+            Metric value on the selected split.
+        group_losses
+            Per-group losses from the model.
+        """
+        if self._solver is None:  # 1) lazy‐instantiate solver if not given
+            self._solver = self.solver_cls(**self.solver_kwargs)
+
+        # 2) bind oracle if user didn’t already
+        if self._solver.oracle is None:
+            self._solver.set_oracle(self.oracle)
+        solver = self._solver
+        solver.run(**run_kwargs)
+        self.group_weights = solver.group_weights
+        print(f"Group weights: {self.group_weights}")
+        self.model.set_group_weights(self.group_weights)
+        self.model.fit(self.X, self.y, self.dset.groups)
+        self.coef_ = self.model.coef_
+        self.intercept_ = self.model.intercept_
+        self.fairness = self.metric.fun(self.model.coef_, self.dset, self.train_test)
+        self.group_losses = self.model._group_loss(self.dset)

badr/__init__.py

@@ -0,0 +1,13 @@
+from . import algorithms, datasets, metrics, models, oracles
+from .__base__ import Badr
+from ._version import __version__
+
+__all__ = [
+    "__version__",
+    "datasets",
+    "metrics",
+    "models",
+    "oracles",
+    "algorithms",
+    "Badr",
+]

badr/_version.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

badr/algorithms/__base__.py

@@ -0,0 +1,90 @@
+from datetime import datetime
+
+import jax.numpy as jnp
+import numpy as np
+
+from badr.oracles import Oracle
+
+
+class Trace:
+    def __init__(self, f=None, freq=1) -> None:
+        self.trace_x = []
+        self.trace_time = []
+        self.trace_fx = []
+        self.start = datetime.now()
+        self._counter = 0
+        self.freq = int(freq)
+        self.f = f
+
+    def __call__(self, dl):
+        if self._counter % self.freq == 0:
+            if self.f is not None:
+                self.trace_fx.append(float(self.f(jnp.array(dl["x"]))))
+            else:
+                self.trace_x.append(np.copy(dl["x"]))
+            delta = (datetime.now() - self.start).total_seconds()
+            self.trace_time.append(delta)
+        self._counter += 1
+
+
+class Algorithm:
+    """
+    Base class for algorithms optimizing group weights.
+
+    Stores common bookkeeping (objective trace, timings, iterates) and a reference
+    to an :class:`~badr.oracles.Oracle`.
+
+    Parameters
+    ----------
+    name : str
+        Display name for the algorithm.
+
+    Attributes
+    ----------
+    oracle : Oracle or None
+        Oracle providing ``fun``/``grad`` (and possibly stochastic primitives).
+    n_groups : int
+        Number of groups (set from the oracle).
+    group_weights : jax.numpy.ndarray
+        Latest group-weight iterate (typically on the simplex).
+    history_f : list[float]
+        Traced objective values (when enabled by the algorithm / ``trace`` flag).
+    history_time : list[float]
+        Elapsed time trace (seconds).
+    history_lambda : list[jax.numpy.ndarray]
+        Traced group-weight iterates (when enabled).
+    """
+
+    def __init__(self, name: str) -> None:
+        self.name = name
+        self.history_f = []
+        self.history_time = []
+        self.history_lambda = []
+        self._last_callback_time = None
+        self.group_weights: jnp.ndarray = jnp.array([])
+        self.oracle = None
+
+    def set_oracle(self, oracle: Oracle) -> None:
+        # Set the oracle for the algorithm.
+        self.oracle = oracle
+        self.n_groups = oracle.n_groups
+
+    def run(self, max_iter: int = 1, verbose: int = 1, trace: bool = False):
+        """
+        Run the algorithm.
+
+        Parameters
+        ----------
+        max_iter : int, default=1
+            Maximum number of iterations.
+        verbose : int, default=1
+            Verbosity level (interpretation is algorithm-specific).
+        trace : bool, default=False
+            If True, record per-iteration history (objective/iterates/times).
+
+        Raises
+        ------
+        NotImplementedError
+            If not implemented by a subclass.
+        """
+        raise NotImplementedError("Algorithm.run() must be implemented in subclasses.")

badr/algorithms/__init__.py

@@ -0,0 +1,13 @@
+from .__base__ import Algorithm
+from .badr import BADRSGD
+from .frank_wolfe import FrankWolfe
+from .slsqp import SLSQP
+from .trust_constr import TrustConstr
+
+__all__ = [
+    "Algorithm",
+    "BADRSGD",
+    "FrankWolfe",
+    "SLSQP",
+    "TrustConstr",
+]

badr/algorithms/badr.py

@@ -0,0 +1,247 @@
+import time
+import jax.numpy as jnp
+from .__base__ import Algorithm
+
+
+class BADRSGD(Algorithm):
+    """
+    Stochastic BADR updates over (w, v, lambda) with simplex projection.
+
+    Uses a :class:`~badr.oracles.StochasticOracle`-style interface:
+    draws a minibatch, forms a group-weighted inner gradient for ``w``, updates an
+    auxiliary vector ``v`` using Hessian-vector products, and updates ``lambda``
+    via a clipped step followed by projection onto the simplex.
+
+    Parameters
+    ----------
+    w0 : jax.numpy.ndarray
+        Initial parameter vector for the lower-level variable ``w``.
+    batch_size : int, default=1
+        Total minibatch size across groups.
+    step_w : float, default=1e-1
+        Step size for ``w`` and ``v`` updates (used as ``step_1``).
+    step_v : float, default=1e-1
+        Stored step size for ``v`` (currently not used in the shown code).
+    step_lambda : float, default=1.0
+        Step size for ``lambda`` update (used as ``step_2``).
+    clip_value : float, default=1.0
+        L2-norm clipping threshold applied to the ``lambda`` gradient estimate.
+
+    Attributes
+    ----------
+    primal_solution : jax.numpy.ndarray or None
+        Final ``w`` iterate.
+    aux_solution : jax.numpy.ndarray or None
+        Final ``v`` iterate.
+    message : str or None
+        Status message after :meth:`run`.
+    history_w, history_v, history_lambda : list[jax.numpy.ndarray]
+        Iterates recorded during :meth:`run`.
+    history_inner_loss_batch : list[float]
+        ``f_hat(w, lambda; batch)`` over iterations.
+    history_outer_metric : list[float]
+        Metric value on the full split at current ``w``.
+    history_norm_grad_w, history_norm_v, history_norm_jt_v, history_norm_grad_H_w : list[float]
+        Diagnostic norms recorded per iteration.
+    history_lambda_entropy : list[float]
+        Entropy ``-sum(lambda log lambda)`` per iteration.
+    history_clip_fraction : list[float]
+        Running fraction of iterations where lambda-gradient clipping activated.
+
+    Notes
+    -----
+    - Simplex projection uses sorting-based Euclidean projection.
+    - Clipping is detected by comparing the pre/post L2 norms of the gradient.
+    """
+
+    def __init__(
+        self,
+        w0: jnp.ndarray,
+        batch_size: int = 1,
+        step_w: float = 1e-1,
+        step_v: float = 1e-1,
+        step_lambda: float = 1.0,
+        clip_value: float = 1.0,
+    ) -> None:
+        super().__init__("BADRSGD")
+        self.w0 = jnp.array(w0)
+        self.batch_size = batch_size
+        self.step_w = step_w
+        self.step_v = step_v
+        self.step_lambda = step_lambda
+        self.clip_value = clip_value
+
+        # existing histories
+        self.history_w = []
+        self.history_v = []
+        self.history_lambda = []
+        self.history_time = []
+        self.message = None
+        self.primal_solution = None
+        self.aux_solution = None
+
+        # NEW: metric histories
+        self.history_inner_loss_batch = []
+        self.history_outer_metric = []
+        self.history_norm_grad_w = []
+        self.history_norm_v = []
+        self.history_norm_jt_v = []
+        self.history_norm_grad_H_w = []
+        self.history_lambda_entropy = []
+        self.history_clip_fraction = []  # running fraction in [0,1]
+
+    @staticmethod
+    def _project_simplex(v: jnp.ndarray, radius: float = 1.0) -> jnp.ndarray:
+        v = jnp.asarray(v)
+        v_sorted = jnp.sort(v)[::-1]
+        cssv = jnp.cumsum(v_sorted) - radius
+        ind = jnp.arange(1, v.shape[0] + 1)
+        cond = v_sorted - cssv / ind > 0
+        rho = max(int(jnp.sum(cond)) - 1, 0)
+        theta = cssv[rho] / (rho + 1)
+        return jnp.maximum(v - theta, 0.0)
+
+    def _clip(self, g: jnp.ndarray) -> jnp.ndarray:
+        norm = jnp.linalg.norm(g)
+        factor = jnp.minimum(1.0, self.clip_value / (norm + 1e-12))
+        return factor * g
+
+    def run(self, max_iter: int = 1000, verbose: int = 1, trace: bool = False):
+        """
+        Run stochastic BADR iterations.
+
+        Parameters
+        ----------
+        max_iter : int, default=1000
+            Number of iterations.
+        verbose : int, default=1
+            If > 0, prints a completion message.
+        trace : bool, default=False
+            If True, also appends the outer metric to ``history_f`` each iteration.
+
+        Returns
+        -------
+        jax.numpy.ndarray
+            Final group weights ``lambda``.
+
+        Raises
+        ------
+        ValueError
+            If no oracle has been set.
+        """
+        if self.oracle is None:
+            raise ValueError("Oracle not set. Please set the oracle before running.")
+
+        w = jnp.array(self.w0)
+        v = jnp.zeros_like(w)
+        lmbda = jnp.ones(self.oracle.n_groups) / self.oracle.n_groups
+
+        # reset histories
+        self.history_w = [w]
+        self.history_v = [v]
+        self.history_lambda = [lmbda]
+        self.history_time = [0.0]
+        if trace:
+            self.history_f = [
+                float(
+                    self.oracle.metric.fun(w, self.oracle.dset, self.oracle.train_test)
+                )
+            ]
+        else:
+            self.history_f = []
+
+        # also reset metric histories
+        self.history_inner_loss_batch = []
+        self.history_outer_metric = []
+        self.history_norm_grad_w = []
+        self.history_norm_v = []
+        self.history_norm_jt_v = []
+        self.history_norm_grad_H_w = []
+        self.history_lambda_entropy = []
+        self.history_clip_fraction = []
+
+        start_t = time.perf_counter()
+        step_1 = self.step_w
+        step_2 = self.step_lambda
+        clipped_count = 0  # running count of clipped λ-updates
+
+        for t in range(max_iter):
+            # --- draw ONE inner batch and reuse it for all inner terms (SOBA)
+            batch, self.oracle.key = self.oracle.sample_batch(
+                self.oracle.key, self.batch_size
+            )
+
+            # 1) inner gradient g_w = ∇_w f_inner(w, λ; batch)
+            grad_F = self.oracle.grad_lower_groups(w, batch)  # shape (S, d)
+            grad_w = grad_F.T @ lmbda  # shape (d,)
+            w_next = w - step_1 * grad_w
+
+            # 2) outer gradient pieces (metric)
+            grad_H_w, grad_H_lambda = self.oracle.grad_upper(w, lmbda)
+
+            # 3) v-update: v ← v - ρ( H_{ww} v + ∇_w H )
+            hvp_weighted = self.oracle.hvp_w_f_hat(w, lmbda, v, batch)  # shape (d,)
+            v_next = v - step_1 * (grad_H_w + hvp_weighted)
+
+            # 4) λ-update: λ ← Proj_Δ( λ - γ * Clip( J^T v + ∇_λ H ) )
+            jt_v = self.oracle.jt_v_lambda_of_grad_w_f_hat(w, v, batch)  # shape (S,)
+            lambda_grad_est = jt_v + grad_H_lambda
+            # detect clipping (norm reduced)
+            pre_norm = float(jnp.linalg.norm(lambda_grad_est))
+            clipped = self._clip(lambda_grad_est)
+            post_norm = float(jnp.linalg.norm(clipped))
+            did_clip = 1.0 if post_norm + 1e-12 < pre_norm else 0.0
+            clipped_count += did_clip
+
+            lambda_candidate = lmbda - step_2 * clipped
+            lambda_next = self._project_simplex(lambda_candidate)
+
+            # step
+            w, v, lmbda = w_next, v_next, lambda_next
+
+            # --- logging ---
+            # inner loss on the sampled batch
+            inner_loss_batch = float(self.oracle.f_hat(w, lmbda, batch))
+            # outer metric (current w on full data; same as your trace metric)
+            outer_metric_val = float(
+                self.oracle.metric.fun(w, self.oracle.dset, self.oracle.train_test)
+            )
+
+            # norms
+            norm_grad_w = float(jnp.linalg.norm(grad_w))
+            norm_v = float(jnp.linalg.norm(v))
+            norm_jt_v = float(jnp.linalg.norm(jt_v))
+            norm_grad_H_w = float(jnp.linalg.norm(grad_H_w))
+
+            # λ entropy
+            lambda_entropy = float(-jnp.sum(lmbda * jnp.log(lmbda + 1e-12)))
+
+            # fraction of clipped λ-updates up to now
+            clip_fraction = clipped_count / float(t + 1)
+
+            # bookkeeping
+            elapsed = time.perf_counter() - start_t
+            self.history_w.append(w)
+            self.history_v.append(v)
+            self.history_lambda.append(lmbda)
+            self.history_time.append(elapsed)
+            if trace:
+                self.history_f.append(outer_metric_val)
+
+            # store metrics
+            self.history_inner_loss_batch.append(inner_loss_batch)
+            self.history_outer_metric.append(outer_metric_val)
+            self.history_norm_grad_w.append(norm_grad_w)
+            self.history_norm_v.append(norm_v)
+            self.history_norm_jt_v.append(norm_jt_v)
+            self.history_norm_grad_H_w.append(norm_grad_H_w)
+            self.history_lambda_entropy.append(lambda_entropy)
+            self.history_clip_fraction.append(clip_fraction)
+
+        self.primal_solution = w
+        self.aux_solution = v
+        self.group_weights = lmbda
+        self.message = f"Completed {max_iter} (stochastic) iterations."
+        if verbose > 0:
+            print("[BADR] Message:", self.message)
+        return self.group_weights

badr/algorithms/frank_wolfe.py

@@ -0,0 +1,137 @@
+from datetime import datetime
+from functools import partial
+
+import jax.numpy as jnp
+from jax import jit
+
+from .__base__ import Algorithm
+
+
+class FrankWolfe(Algorithm):
+    """
+    Frank--Wolfe on the simplex using a linear minimization oracle.
+
+    Uses the oracle gradient ``g = oracle.grad(x)`` and the simplex LMO that
+    returns the vertex at ``argmin_i g_i``. The update uses a diminishing step
+    size given in :meth:`_step`. Convergence is checked with the FW gap
+    ``g_t = -g^T (x - x_prev)``.
+
+    Parameters
+    ----------
+    starting_point : jax.numpy.ndarray or None, optional
+        Initial point on the simplex. If None, uses the uniform distribution.
+    eps : float, default=1e-6
+        Stopping threshold on the FW gap.
+
+    Attributes
+    ----------
+    iterates : list[jax.numpy.ndarray]
+        Stored iterates (always filled).
+    success : bool
+        True if the stopping condition was reached.
+    message : str or None
+        Status message after :meth:`run`.
+    """
+
+    def __init__(self, starting_point=None, eps=1e-6):
+        super().__init__("Frank-Wolfe")
+        self.oracle = None
+        self.eps = eps
+        self.starting_point = starting_point
+        self.success = False
+        self.iterates = []
+        self.history_x = []
+        self.history_time = []
+        self.history_fx = []
+        self.message = None
+
+    @staticmethod
+    def _lmo(grad: jnp.ndarray) -> jnp.ndarray:
+        v = jnp.zeros_like(grad)
+        idx = jnp.argmin(grad)
+        return v.at[idx].set(1.0)
+
+    @staticmethod
+    @partial(jit, static_argnums=(2,))
+    def _step(x: jnp.ndarray, g: jnp.ndarray, iteration: int) -> jnp.ndarray:
+        v = FrankWolfe._lmo(g)
+        d = v - x
+        alpha = jnp.minimum(
+            1.0,
+            (2.0 + jnp.log(iteration + 1)) / (iteration + 2 + jnp.log(iteration + 1)),
+        )
+        return x + alpha * d
+
+    def run(self, max_iter: int = 300, verbose: int = 1, trace: bool = False):
+        """
+        Run Frank--Wolfe iterations.
+
+        Parameters
+        ----------
+        max_iter : int, default=300
+            Maximum number of iterations.
+        verbose : int, default=1
+            If > 0, prints success and message.
+        trace : bool, default=False
+            If True, records iterates and elapsed time.
+
+        Returns
+        -------
+        jax.numpy.ndarray
+            Final group weights.
+
+        Raises
+        ------
+        ValueError
+            If no oracle has been set.
+        """
+        self.group_weights = jnp.zeros(self.n_groups)
+        if self.starting_point is None:
+            self.starting_point = jnp.full(self.n_groups, 1.0 / self.n_groups)
+        else:
+            self.starting_point = self.starting_point
+        if self.oracle is None:
+            raise ValueError("Oracle not set. Please set the oracle before running.")
+        if trace:
+            self.start = datetime.now()
+
+        x = self.starting_point
+        self.iterates = [x]
+
+        for it in range(max_iter):
+            g = self.oracle.grad(x)
+
+            x = FrankWolfe._step(x, g, it)
+
+            self.iterates.append(x)
+
+            if trace:
+                now = datetime.now()
+                self.history_time.append((now - self.start).total_seconds())
+                self.history_x.append(x)
+
+            g_t = jnp.dot(-g, x - self.iterates[-2])
+            if g_t < self.eps:
+                self.success = True
+                break
+
+        self.group_weights = x
+        self.message = f"Converged in {len(self.iterates) - 1} iterations (over {max_iter} iterations)."
+        if verbose > 0:
+            print(f"[FW] Success: {self.success}")
+            print(f"[FW] Message: {self.message}")
+        return self.group_weights
+
+    def postprocess(self):
+        """
+        Populate ``history_f`` from the stored ``history_x``.
+
+        Raises
+        ------
+        ValueError
+            If no oracle has been set.
+        """
+        if self.oracle is None:
+            raise ValueError("Oracle not set. Please set the oracle before running.")
+        for x in self.history_x:
+            self.history_f.append(self.oracle.fun(x))