AQR Data Scientist Interview Guide

import numpy as np


def factor_neutral_weights(mu, B, lower, upper, tol=1e-9, max_iter=20000):
    """Return weights w maximizing mu^T w with sum(w)=0 and B^T w=0 under box bounds.

    Solves the linear program via a feasibility-preserving active-set style method:
    maintain a set of free variables, enforce equalities on free variables exactly,
    then move along the nullspace direction that increases objective until a bound
    hits. If constraints are infeasible, raises ValueError.
    """
    mu = np.asarray(mu, dtype=float).reshape(-1)
    lower = np.asarray(lower, dtype=float).reshape(-1)
    upper = np.asarray(upper, dtype=float).reshape(-1)

    if mu.ndim != 1:
        raise ValueError("mu must be 1D")
    n = mu.shape[0]
    if lower.shape[0] != n or upper.shape[0] != n:
        raise ValueError("lower/upper must have length n")
    if np.any(lower > upper + tol):
        raise ValueError("Infeasible bounds: lower > upper")

    B = np.asarray(B, dtype=float)
    if B.ndim == 1:
        B = B.reshape(n, 1)
    if B.shape[0] != n:
        raise ValueError("B must have shape (n,k)")
    k = B.shape[1]

    # Equality constraints A^T w = 0 where A is (n, m) with m = 1 + k
    A = np.column_stack([np.ones(n), B])
    m = A.shape[1]

    # Start at w=0 if within bounds, else project to middle and check feasibility later.
    w = np.zeros(n)
    if np.any(w < lower - tol) or np.any(w > upper + tol):
        w = np.clip(w, lower, upper)

    # Free set: indices not fixed at bounds.
    def is_at_lower(i):
        return w[i] <= lower[i] + 10 * tol

    def is_at_upper(i):
        return w[i] >= upper[i] - 10 * tol

    free = np.array([not (is_at_lower(i) or is_at_upper(i)) for i in range(n)], dtype=bool)

    def solve_for_free(w_fixed, free_mask):
        """Find w_free so that A^T w = 0, with fixed entries held constant.

        Returns a feasible w if solvable, else None.
        """
        idx_f = np.where(free_mask)[0]
        idx_c = np.where(~free_mask)[0]

        Af = A[idx_f, :]  # (nf, m)
        Ac = A[idx_c, :]
        rhs = -Ac.T @ w_fixed[idx_c]  # (m,)

        # Need Af.T w_f = rhs.
        # Choose minimum-norm solution if underdetermined.
        # Check consistency via least squares residual.
        if idx_f.size == 0:
            if np.linalg.norm(rhs, ord=np.inf) <= tol:
                return w_fixed.copy()
            return None

        # Solve Af.T w_f = rhs.
        M = Af.T  # (m, nf)
        w_f, residuals, rank, s = np.linalg.lstsq(M, rhs, rcond=None)

        if np.linalg.norm(M @ w_f - rhs, ord=np.inf) > 100 * tol:
            return None

        w_new = w_fixed.copy()
        w_new[idx_f] = w_f
        return w_new

    # First, attempt to find any feasible point by iteratively freeing variables.
    for _ in range(n + 5):
        w_try = solve_for_free(w, free)
        if w_try is None:
            # Not enough degrees of freedom, free one more variable.
            # Prefer to free the variable with widest interval.
            idx_c = np.where(~free)[0]
            if idx_c.size == 0:
                break
            widths = upper[idx_c] - lower[idx_c]
            j = idx_c[int(np.argmax(widths))]
            free[j] = True
            continue
        # Check bounds.
        if np.all(w_try >= lower - 1e-7) and np.all(w_try <= upper + 1e-7):
            w = np.clip(w_try, lower, upper)
            break
        # If out of bounds, clamp offending variables and mark as fixed.
        viol_low = w_try < lower - 1e-7
        viol_up = w_try > upper + 1e-7
        if not (np.any(viol_low) or np.any(viol_up)):
            w = np.clip(w_try, lower, upper)
            break
        for i in np.where(viol_low)[0]:
            w[i] = lower[i]
            free[i] = False
        for i in np.where(viol_up)[0]:
            w[i] = upper[i]
            free[i] = False
    else:
        raise ValueError("Infeasible constraints")

    # Verify feasibility.
    if np.linalg.norm(A.T @ w, ord=np.inf) > 1e-6:
        # One more attempt: make all non-bound variables free and solve.
        free = np.array([not (is_at_lower(i) or is_at_upper(i)) for i in range(n)], dtype=bool)
        w_try = solve_for_free(w, free)
        if w_try is None or np.any(w_try < lower - 1e-6) or np.any(w_try > upper + 1e-6):
            raise ValueError("Infeasible constraints")
        w = np.clip(w_try, lower, upper)

    # Active-set improvement iterations.
    for _ in range(max_iter):
        idx_f = np.where(free)[0]
        idx_c = np.where(~free)[0]

        # Compute a nullspace direction d with A^T d = 0, supported on free vars.
        if idx_f.size == 0:
            break

        Af = A[idx_f, :]           # (nf, m)
        mu_f = mu[idx_f]

        # We want d_f in Null(Af.T) maximizing mu_f^T d_f, normalized.
        # Project mu_f onto nullspace of Af.T.
        M = Af.T                   # (m, nf)
        # Orthonormal basis for rowspace of M: Q from QR of M.T gives colspace(M.T)=rowspace(M)
        # Compute projection onto nullspace via P = I - V V^T where V spans rowspace(M).
        try:
            Q, _ = np.linalg.qr(M.T, mode='reduced')  # Q: (nf, r)
        except np.linalg.LinAlgError:
            Q = None

        if Q is None or Q.size == 0:
            d_f = mu_f.copy()
        else:
            d_f = mu_f - Q @ (Q.T @ mu_f)

        norm = np.linalg.norm(d_f)
        if norm <= 1e-12:
            # No improving direction in nullspace.
            break
        d_f = d_f / norm

        d = np.zeros(n)
        d[idx_f] = d_f

        # Step until a bound hits.
        alpha_max = np.inf
        hit = None
        for i in idx_f:
            if d[i] > 1e-14:
                a = (upper[i] - w[i]) / d[i]
                if a < alpha_max:
                    alpha_max = a
                    hit = (i, 'upper')
            elif d[i] < -1e-14:
                a = (lower[i] - w[i]) / d[i]
                if a < alpha_max:
                    alpha_max = a
                    hit = (i, 'lower')

        if not np.isfinite(alpha_max) or alpha_max <= 1e-14:
            break

        # Take step.
        w_new = w + alpha_max * d
        w_new = np.clip(w_new, lower, upper)

        # Enforce equalities again by solving for free variables.
        # Fix the hit variable at its bound to maintain feasibility.
        if hit is not None:
            i_hit, which = hit
            w_new[i_hit] = upper[i_hit] if which == 'upper' else lower[i_hit]
            free_new = free.copy()
            free_new[i_hit] = False
        else:
            free_new = free.copy()

        w_feas = solve_for_free(w_new, free_new)
        if w_feas is None:
            # If re-solving failed, fix one more variable with smallest slack.
            slacks = np.minimum(w_new - lower, upper - w_new)
            candidates = np.where(free_new)[0]
            if candidates.size == 0:
                break
            j = candidates[int(np.argmin(slacks[candidates]))]
            # Fix to nearest bound.
            if (w_new[j] - lower[j]) <= (upper[j] - w_new[j]):
                w_new[j] = lower[j]
            else:
                w_new[j] = upper[j]
            free_new[j] = False
            w_feas = solve_for_free(w_new, free_new)
            if w_feas is None:
                raise ValueError("Infeasible constraints")

        if np.any(w_feas < lower - 1e-6) or np.any(w_feas > upper + 1e-6):
            # Clamp and continue with updated active set.
            w = np.clip(w_feas, lower, upper)
        else:
            w = w_feas

        free = np.array([not (w[i] <= lower[i] + 10 * tol or w[i] >= upper[i] - 10 * tol) for i in range(n)], dtype=bool)

        # Check feasibility.
        if np.linalg.norm(A.T @ w, ord=np.inf) > 1e-6:
            w_try = solve_for_free(w, free)
            if w_try is None:
                raise ValueError("Infeasible constraints")
            if np.any(w_try < lower - 1e-6) or np.any(w_try > upper + 1e-6):
                raise ValueError("Infeasible constraints")
            w = w_try

    # Final feasibility check.
    if np.linalg.norm(A.T @ w, ord=np.inf) > 1e-5:
        raise ValueError("Infeasible constraints")
    if np.any(w < lower - 1e-6) or np.any(w > upper + 1e-6):
        raise ValueError("Infeasible constraints")

    return w