Kernel mean matching

neureptrace.decoding.kernel_mean_matching implements dependency-light kernel mean matching (KMM) for unlabeled target-adaptive source weighting.

KMM estimates non-negative source-row weights so the weighted source feature distribution better matches an unlabeled target feature distribution in a kernel space. This is useful for covariate-shift correction in cross-subject decoding.

Protocol interpretation:

• uses source features X_s,
• optionally uses source labels y_s for class-balanced post-normalization,
• uses unlabeled target features X_t,
• does not accept held-out target labels y_t.

Therefore, KMM is a Protocol 2 / unlabeled target-adaptive method.

Typical use:

from neureptrace.decoding.kernel_mean_matching import kernel_mean_matching_weights

result = kernel_mean_matching_weights(
    X_source,
    X_target_unlabeled,
    kernel="rbf",
    gamma="median",
    max_weight=10.0,
)

sample_weight = result.weights

neureptrace.decoding.kernel_mean_matching

Kernel mean matching for unlabeled target-adaptive source weighting.

Kernel mean matching (KMM) reweights source rows so the weighted source feature distribution better matches an unlabeled held-out target feature distribution in an RKHS. The implementation here is deliberately dependency-light: it uses NumPy plus SciPy's SLSQP optimizer rather than an external quadratic-programming package.

The public API is Category 2 in the cross-subject protocol taxonomy. It uses source features and unlabeled target features, and may optionally use source labels only for class-balanced post-normalization. It does not accept held-out target labels.

KernelMeanMatchingConfig dataclass

Configuration for dependency-light KMM source weighting.

Source code in src/neureptrace/decoding/kernel_mean_matching.py

@dataclass(frozen=True, slots=True)
class KernelMeanMatchingConfig:
    """Configuration for dependency-light KMM source weighting."""

    kernel: str = "rbf"
    gamma: float | str = DEFAULT_KMM_GAMMA
    max_weight: float = DEFAULT_KMM_MAX_WEIGHT
    epsilon: float | str | None = DEFAULT_KMM_EPSILON
    regularization: float = DEFAULT_KMM_REGULARIZATION
    max_iter: int = DEFAULT_KMM_MAX_ITER
    tol: float = DEFAULT_KMM_TOL
    normalize: bool = True
    class_balance: bool = False

KernelMeanMatchingResult dataclass

KMM source weights and provenance metadata.

Source code in src/neureptrace/decoding/kernel_mean_matching.py

@dataclass(frozen=True, slots=True)
class KernelMeanMatchingResult:
    """KMM source weights and provenance metadata."""

    weights: np.ndarray
    objective_value: float
    success: bool
    status: str
    metadata: dict[str, Any] = field(default_factory=dict)

kernel_mean_matching_weights(source_features, target_features, *, kernel='rbf', gamma=DEFAULT_KMM_GAMMA, max_weight=DEFAULT_KMM_MAX_WEIGHT, epsilon=DEFAULT_KMM_EPSILON, regularization=DEFAULT_KMM_REGULARIZATION, max_iter=DEFAULT_KMM_MAX_ITER, tol=DEFAULT_KMM_TOL, normalize=True, source_labels=None, class_balance=False)

Estimate source-row weights from unlabeled target features.

Parameters:

Name	Type	Description	Default
source_features	Sequence[Sequence[float]] | ndarray	Source feature matrix, usually pooled over source subjects in an outer LOSO fold.	required
target_features	Sequence[Sequence[float]] | ndarray	Unlabeled held-out target feature matrix used for covariate-shift weighting. The function intentionally has no target_labels argument.	required
kernel	str | None	"rbf" or "linear".	'rbf'
gamma	float | str	RBF kernel width. "median"/"auto" uses the median pairwise squared distance over source and target rows. "scale" uses 1 / n_features.	DEFAULT_KMM_GAMMA
max_weight	float | str	Upper bound for each source weight before optional normalization.	DEFAULT_KMM_MAX_WEIGHT
epsilon	float | str | None	KMM sum-constraint slack. "auto" uses the common (sqrt(n_source)-1)/sqrt(n_source) heuristic. None disables the sum constraint and only enforces box bounds.	DEFAULT_KMM_EPSILON
regularization	float | str	Diagonal ridge added to the source-source kernel matrix.	DEFAULT_KMM_REGULARIZATION
max_iter	int | str	SLSQP optimizer controls.	DEFAULT_KMM_MAX_ITER
tol	int | str	SLSQP optimizer controls.	DEFAULT_KMM_MAX_ITER
normalize	bool	If true, final weights are rescaled to have mean one.	True
source_labels	Sequence[Any] | ndarray | None	Optional source labels used only when class_balance=True.	None
class_balance	bool	If true, rescale weights so each source class receives equal total mass. This uses source labels only and remains Category 2.	False

Returns:

Type	Description
KernelMeanMatchingResult	Mean-one source sample weights and protocol metadata.

Source code in src/neureptrace/decoding/kernel_mean_matching.py

def kernel_mean_matching_weights(
    source_features: Sequence[Sequence[float]] | np.ndarray,
    target_features: Sequence[Sequence[float]] | np.ndarray,
    *,
    kernel: str | None = "rbf",
    gamma: float | str = DEFAULT_KMM_GAMMA,
    max_weight: float | str = DEFAULT_KMM_MAX_WEIGHT,
    epsilon: float | str | None = DEFAULT_KMM_EPSILON,
    regularization: float | str = DEFAULT_KMM_REGULARIZATION,
    max_iter: int | str = DEFAULT_KMM_MAX_ITER,
    tol: float | str = DEFAULT_KMM_TOL,
    normalize: bool = True,
    source_labels: Sequence[Any] | np.ndarray | None = None,
    class_balance: bool = False,
) -> KernelMeanMatchingResult:
    """Estimate source-row weights from unlabeled target features.

    Parameters
    ----------
    source_features:
        Source feature matrix, usually pooled over source subjects in an outer
        LOSO fold.
    target_features:
        Unlabeled held-out target feature matrix used for covariate-shift
        weighting.  The function intentionally has no ``target_labels`` argument.
    kernel:
        ``"rbf"`` or ``"linear"``.
    gamma:
        RBF kernel width.  ``"median"``/``"auto"`` uses the median pairwise
        squared distance over source and target rows.  ``"scale"`` uses
        ``1 / n_features``.
    max_weight:
        Upper bound for each source weight before optional normalization.
    epsilon:
        KMM sum-constraint slack.  ``"auto"`` uses the common
        ``(sqrt(n_source)-1)/sqrt(n_source)`` heuristic.  ``None`` disables the
        sum constraint and only enforces box bounds.
    regularization:
        Diagonal ridge added to the source-source kernel matrix.
    max_iter, tol:
        SLSQP optimizer controls.
    normalize:
        If true, final weights are rescaled to have mean one.
    source_labels:
        Optional source labels used only when ``class_balance=True``.
    class_balance:
        If true, rescale weights so each source class receives equal total mass.
        This uses source labels only and remains Category 2.

    Returns
    -------
    KernelMeanMatchingResult
        Mean-one source sample weights and protocol metadata.
    """

    source = _feature_matrix(source_features, name="source_features")
    target = _feature_matrix(target_features, name="target_features")
    if source.shape[1] != target.shape[1]:
        raise ValueError(f"source_features and target_features must have the same feature width: {source.shape[1]} != {target.shape[1]}.")
    normalized_kernel = normalize_kmm_kernel(kernel)
    max_weight_value = _positive_float(max_weight, name="max_weight")
    regularization_value = _nonnegative_float(regularization, name="regularization")
    max_iterations = _positive_int(max_iter, name="max_iter")
    tolerance = _positive_float(tol, name="tol")
    epsilon_value = normalize_kmm_epsilon(epsilon, n_source=source.shape[0])
    gamma_value = resolve_kmm_gamma(gamma, source, target) if normalized_kernel == "rbf" else ""

    source_kernel = _source_kernel(source, normalized_kernel, gamma_value)
    if regularization_value > 0.0:
        source_kernel = source_kernel + regularization_value * np.eye(source.shape[0], dtype=float)
    kappa = _kappa_vector(source, target, normalized_kernel, gamma_value)
    initial = np.ones(source.shape[0], dtype=float)
    bounds = [(0.0, max_weight_value) for _ in range(source.shape[0])]
    constraints = _sum_constraints(source.shape[0], epsilon_value)

    def objective(weights: np.ndarray) -> float:
        return float(0.5 * weights @ source_kernel @ weights - kappa @ weights)

    def gradient(weights: np.ndarray) -> np.ndarray:
        return source_kernel @ weights - kappa

    solution = minimize(
        objective,
        initial,
        jac=gradient,
        method="SLSQP",
        bounds=bounds,
        constraints=constraints,
        options={"maxiter": max_iterations, "ftol": tolerance, "disp": False},
    )
    raw_weights = np.asarray(solution.x, dtype=float) if np.all(np.isfinite(solution.x)) else initial.copy()
    raw_weights = _project_weights(raw_weights, max_weight=max_weight_value, epsilon=epsilon_value, n_source=source.shape[0])
    status = str(solution.message)
    success = bool(solution.success)

    if not success:
        fallback = minimize(
            objective,
            initial,
            jac=gradient,
            method="L-BFGS-B",
            bounds=bounds,
            options={"maxiter": max_iterations, "ftol": tolerance},
        )
        if np.all(np.isfinite(fallback.x)) and objective(np.asarray(fallback.x, dtype=float)) <= objective(raw_weights) + 1.0e-10:
            raw_weights = _project_weights(np.asarray(fallback.x, dtype=float), max_weight=max_weight_value, epsilon=epsilon_value, n_source=source.shape[0])
            status = f"SLSQP failed ({status}); used bounded L-BFGS-B fallback: {fallback.message}"
            success = bool(fallback.success)

    weights = raw_weights.copy()
    if class_balance:
        if source_labels is None:
            raise ValueError("source_labels are required when class_balance=True.")
        weights = _class_balanced_weights(weights, source_labels)
    if normalize:
        weights = _mean_one_weights(weights)
    objective_value = objective(raw_weights)
    metadata = _metadata(
        n_source=source.shape[0],
        n_target=target.shape[0],
        feature_dim=source.shape[1],
        kernel=normalized_kernel,
        gamma=gamma_value,
        max_weight=max_weight_value,
        epsilon=epsilon_value,
        regularization=regularization_value,
        max_iter=max_iterations,
        tol=tolerance,
        normalize=bool(normalize),
        class_balance=bool(class_balance),
        success=success,
        status=status,
        objective_value=objective_value,
        weights=weights,
    )
    return KernelMeanMatchingResult(
        weights=weights.astype(np.float64, copy=False),
        objective_value=float(objective_value),
        success=success,
        status=status,
        metadata=metadata,
    )

kernel_mean_matching_weights_from_config(source_features, target_features, config=None, *, source_labels=None)

Estimate KMM weights from a dataclass or mapping configuration.

Source code in src/neureptrace/decoding/kernel_mean_matching.py

def kernel_mean_matching_weights_from_config(
    source_features: Sequence[Sequence[float]] | np.ndarray,
    target_features: Sequence[Sequence[float]] | np.ndarray,
    config: KernelMeanMatchingConfig | Mapping[str, Any] | None = None,
    *,
    source_labels: Sequence[Any] | np.ndarray | None = None,
) -> KernelMeanMatchingResult:
    """Estimate KMM weights from a dataclass or mapping configuration."""

    cfg = kmm_config() if config is None else config
    if isinstance(cfg, Mapping):
        cfg = kmm_config(**dict(cfg))
    return kernel_mean_matching_weights(
        source_features,
        target_features,
        kernel=cfg.kernel,
        gamma=cfg.gamma,
        max_weight=cfg.max_weight,
        epsilon=cfg.epsilon,
        regularization=cfg.regularization,
        max_iter=cfg.max_iter,
        tol=cfg.tol,
        normalize=cfg.normalize,
        source_labels=source_labels,
        class_balance=cfg.class_balance,
    )

kmm_config(*, kernel='rbf', gamma=DEFAULT_KMM_GAMMA, max_weight=DEFAULT_KMM_MAX_WEIGHT, epsilon=DEFAULT_KMM_EPSILON, regularization=DEFAULT_KMM_REGULARIZATION, max_iter=DEFAULT_KMM_MAX_ITER, tol=DEFAULT_KMM_TOL, normalize=True, class_balance=False)

Normalize user-facing KMM configuration values.

Source code in src/neureptrace/decoding/kernel_mean_matching.py

def kmm_config(
    *,
    kernel: str | None = "rbf",
    gamma: float | str = DEFAULT_KMM_GAMMA,
    max_weight: float | str = DEFAULT_KMM_MAX_WEIGHT,
    epsilon: float | str | None = DEFAULT_KMM_EPSILON,
    regularization: float | str = DEFAULT_KMM_REGULARIZATION,
    max_iter: int | str = DEFAULT_KMM_MAX_ITER,
    tol: float | str = DEFAULT_KMM_TOL,
    normalize: bool = True,
    class_balance: bool = False,
) -> KernelMeanMatchingConfig:
    """Normalize user-facing KMM configuration values."""

    _reject_boolean_numeric(gamma, error_message=_KMM_GAMMA_ERROR)
    _reject_boolean_numeric(epsilon, error_message=_KMM_EPSILON_ERROR)
    return KernelMeanMatchingConfig(
        kernel=normalize_kmm_kernel(kernel),
        gamma=gamma,
        max_weight=_positive_float(max_weight, name="max_weight"),
        epsilon=epsilon,
        regularization=_nonnegative_float(regularization, name="regularization"),
        max_iter=_positive_int(max_iter, name="max_iter"),
        tol=_positive_float(tol, name="tol"),
        normalize=bool(normalize),
        class_balance=bool(class_balance),
    )

normalize_kmm_kernel(value)

Normalize supported KMM kernel aliases.

Source code in src/neureptrace/decoding/kernel_mean_matching.py

def normalize_kmm_kernel(value: str | None) -> str:
    """Normalize supported KMM kernel aliases."""

    normalized = "rbf" if value is None else str(value).strip().lower().replace("-", "_")
    normalized = {
        "gaussian": "rbf",
        "gaussian_rbf": "rbf",
        "radial_basis": "rbf",
        "dot": "linear",
        "linear_kernel": "linear",
    }.get(normalized, normalized)
    if normalized not in KMM_KERNELS:
        raise ValueError(f"Unknown KMM kernel {value!r}. Available kernels: {', '.join(KMM_KERNELS)}.")
    return normalized

normalize_kmm_epsilon(value, *, n_source)

Resolve KMM sum-constraint slack.

Source code in src/neureptrace/decoding/kernel_mean_matching.py

def normalize_kmm_epsilon(value: float | str | None, *, n_source: int) -> float | None:
    """Resolve KMM sum-constraint slack."""

    _reject_boolean_numeric(value, error_message=_KMM_EPSILON_ERROR)
    if value is None:
        return None
    if isinstance(value, str):
        normalized = value.strip().lower().replace("-", "_")
        if normalized in {"", "none", "off", "disabled"}:
            return None
        if normalized in {"auto", "default"}:
            n = max(1, int(n_source))
            return float((np.sqrt(n) - 1.0) / np.sqrt(n))
        numeric = float(normalized)
    else:
        numeric = float(value)
    if not np.isfinite(numeric) or numeric < 0.0:
        raise ValueError(_KMM_EPSILON_ERROR)
    return numeric

resolve_kmm_gamma(value, source_features, target_features)

Resolve an RBF gamma value from a numeric value or heuristic alias.

Source code in src/neureptrace/decoding/kernel_mean_matching.py

def resolve_kmm_gamma(value: float | str, source_features: Sequence[Sequence[float]] | np.ndarray, target_features: Sequence[Sequence[float]] | np.ndarray) -> float:
    """Resolve an RBF gamma value from a numeric value or heuristic alias."""

    source = _feature_matrix(source_features, name="source_features")
    target = _feature_matrix(target_features, name="target_features")
    if source.shape[1] != target.shape[1]:
        raise ValueError("source_features and target_features must have matching feature width.")
    _reject_boolean_numeric(value, error_message=_KMM_GAMMA_ERROR)
    if isinstance(value, str):
        normalized = value.strip().lower().replace("-", "_")
        if normalized in {"median", "auto", "median_distance", "median_heuristic"}:
            squared = _upper_pairwise_squared_distances(np.vstack([source, target]))
            positive = squared[squared > _MIN_SCALE]
            sigma2 = float(np.median(positive)) if positive.size else 1.0
            return 1.0 / (2.0 * max(sigma2, _MIN_SCALE))
        if normalized == "scale":
            return 1.0 / max(1, source.shape[1])
        numeric = float(normalized)
    else:
        numeric = float(value)
    if not np.isfinite(numeric) or numeric <= 0.0:
        raise ValueError(_KMM_GAMMA_ERROR)
    return numeric