bwd_random

Balancing Walk Design with reversion to Bernoulli randomization.

BWDRandom(N, D, delta=0.05, q=0.5, intercept=True, phi=1)

Bases: object

The Balancing Walk Design with Reversion to Bernoulli Randomization

This is an algorithm from Arbour et al (2022). At each step, it adjusts randomization probabilities to ensure that imbalance tends towards zero. In particular, if current imbalance is w and the current covariate profile is \(x\), then the probability of treatment conditional on history will be:

\[p_i = q \left(1 - \phi \frac{x \cdot w}{\alpha}\right)\]

\(q\) is the desired marginal probability, \(\phi\) is the parameter which controls robustness and \(\alpha\) is the normalizing constant which ensures the probability is well-formed.

If \(|x \cdot w| > \alpha\)

All future units will be assigned by complete randomization.

N : int Total number of points D : int Dimension of the data delta : float, optional Probability of failure, by default 0.05 q : float, optional Target marginal probability of treatment, by default 0.5 intercept : bool, optional Whether an intercept term be added to covariate profiles, by default True phi : float, optional Robustness parameter. A value of 1 focuses entirely on balance, while a value approaching zero does pure randomization, by default 1

Source code in src/bwd/bwd_random.py

def __init__(
    self,
    N: int,
    D: int,
    delta: float = 0.05,
    q: float = 0.5,
    intercept: bool = True,
    phi: float = 1,
) -> None:
    """
    Parameters
    ----------
    N : int
        Total number of points
    D : int
        Dimension of the data
    delta : float, optional
        Probability of failure, by default 0.05
    q : float, optional
        Target marginal probability of treatment, by default 0.5
    intercept : bool, optional
        Whether an intercept term be added to covariate profiles, by default True
    phi : float, optional
        Robustness parameter. A value of 1 focuses entirely on balance, while a value
        approaching zero does pure randomization, by default 1
    """
    self.q = q
    self.intercept = intercept
    self.delta = delta
    self.N = N
    self.D = D + int(self.intercept)

    self.value_plus = 2 * (1 - self.q)
    self.value_minus = -2 * self.q
    self.phi = phi
    self.reset()

definition property

Get the definition parameters of the balancer

Returns:

Type	Description
dict	Dictionary containing N, D, delta, q, intercept, and phi

state property

Get the current state of the balancer

Returns:

Type	Description
dict	Dictionary containing w_i and iterations

assign_all(X)

Assign all points

This assigns units to treatment in the offline setting in which all covariate profiles are available prior to assignment. The algorithm assigns as if units were still only observed in a stream.

Parameters:

Name	Type	Description	Default
X	ndarray	Array of size n × d of covariate profiles	required

Returns:

Type	Description
ndarray	Array of treatment assignments

Source code in src/bwd/bwd_random.py

def assign_all(self, X: np.ndarray) -> np.ndarray:
    """Assign all points

    This assigns units to treatment in the offline setting in which all covariate
    profiles are available prior to assignment. The algorithm assigns as if units
    were still only observed in a stream.

    Parameters
    ----------
    X : np.ndarray
        Array of size n × d of covariate profiles

    Returns
    -------
    np.ndarray
        Array of treatment assignments
    """
    return np.array([self.assign_next(X[i, :]) for i in range(X.shape[0])])

assign_next(x)

Assign treatment to the next point

Parameters:

Name	Type	Description	Default
x	ndarray	Covariate profile of unit to assign treatment	required

Returns:

Type	Description
ndarray	Treatment assignment (0 or 1)

Source code in src/bwd/bwd_random.py

def assign_next(self, x: np.ndarray) -> int:
    """Assign treatment to the next point

    Parameters
    ----------
    x : np.ndarray
        Covariate profile of unit to assign treatment

    Returns
    -------
    np.ndarray
        Treatment assignment (0 or 1)
    """
    if self.intercept:
        x = np.concatenate(([1], x))
    dot = x @ self.w_i
    if abs(dot) > self.alpha:
        self.w_i = np.zeros((self.D,))
        self.set_alpha(self.N - self.iterations)
        dot = 0.0

    p_i = self.q * (1 - self.phi * dot / self.alpha)

    if np.random.rand() < p_i:
        value = self.value_plus
        assignment = 1
    else:
        value = self.value_minus
        assignment = -1
    self.w_i += value * x
    self.iterations += 1
    return int((assignment + 1) / 2)

reset()

Reset the balancer to initial state

Resets the imbalance vector to zeros, initializes alpha, and sets iterations to 0.

Source code in src/bwd/bwd_random.py

def reset(self):
    """Reset the balancer to initial state

    Resets the imbalance vector to zeros, initializes alpha, and sets iterations to 0.
    """
    self.w_i = np.zeros((self.D,))
    self.alpha = np.log(2 * self.N / self.delta) * min(1 / self.q, 9.32)
    self.iterations = 0

set_alpha(N)

Set normalizing constant for remaining N units

Parameters:

Name	Type	Description	Default
N	int	Number of units remaining in the sample	required

Raises:

Type	Description
SampleSizeExpendedError	If N is negative

Source code in src/bwd/bwd_random.py

def set_alpha(self, N: int) -> None:
    """Set normalizing constant for remaining N units

    Parameters
    ----------
    N : int
        Number of units remaining in the sample

    Raises
    ------
    SampleSizeExpendedError
        If N is negative
    """
    if N < 0:
        raise SampleSizeExpendedError()
    self.alpha = -1

update_state(w_i, iterations)

Update the state of the balancer

Parameters:

Name	Type	Description	Default
w_i	array - like	Current imbalance vector	required
iterations	int	Current iteration count	required

Source code in src/bwd/bwd_random.py

def update_state(self, w_i, iterations):
    """Update the state of the balancer

    Parameters
    ----------
    w_i : array-like
        Current imbalance vector
    iterations : int
        Current iteration count
    """
    self.w_i = np.array(w_i)
    self.iterations = iterations