Joint distribution adaptation

neureptrace.decoding.joint_distribution_adaptation implements iterative Category-2 alignment of marginal and class-conditional source-target distributions.

The method uses source labels and unlabeled target features. Target class structure is represented by pseudo-labels or optional source-model target probabilities. Held-out target labels are not part of the API.

neureptrace.decoding.joint_distribution_adaptation

Iterative Joint Distribution Adaptation for cross-subject feature transfer.

The implementation is intentionally protocol-explicit. It uses labeled source features and unlabeled target features. Target class structure is represented by pseudo-labels or optional source-model target probabilities; held-out target labels are not accepted by the public API.

JointDistributionAdaptationConfig dataclass

Configuration for iterative JDA.

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

@dataclass(frozen=True, slots=True)
class JointDistributionAdaptationConfig:
    """Configuration for iterative JDA."""

    method: str = "jda"
    n_components: int | str = 16
    max_iterations: int = 10
    conditional_weight: float = 1.0
    regularization: float = 1e-3
    eigen_ridge: float = 1e-6
    temperature: float = 1.0
    standardize: bool = True
    normalize_latent: bool = False

JointDistributionAdaptationResult dataclass

Projected source/target rows and pseudo-label diagnostics.

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

@dataclass(frozen=True, slots=True)
class JointDistributionAdaptationResult:
    """Projected source/target rows and pseudo-label diagnostics."""

    source_features: np.ndarray
    target_features: np.ndarray
    projection: np.ndarray
    feature_mean: np.ndarray
    feature_scale: np.ndarray
    eigenvalues: np.ndarray
    target_pseudo_labels: np.ndarray
    target_probabilities: np.ndarray
    classes: tuple[Any, ...]
    n_iterations: int
    converged: bool
    metadata: dict[str, Any] = field(default_factory=dict)

fit_joint_distribution_adaptation(source_features, source_labels, target_features, *, target_probabilities=None, classes=None, config=None, method=None, n_components=None, max_iterations=None, conditional_weight=None, regularization=None, eigen_ridge=None, temperature=None, standardize=None, normalize_latent=None)

Fit iterative marginal-plus-conditional source-target alignment.

target_probabilities may contain source-model probabilities for the unlabeled target rows. If omitted, target pseudo-labels are initialized by nearest source-class centroids. No target-label argument exists.

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

def fit_joint_distribution_adaptation(
    source_features: Sequence[Sequence[float]] | np.ndarray,
    source_labels: Sequence[Any] | np.ndarray,
    target_features: Sequence[Sequence[float]] | np.ndarray,
    *,
    target_probabilities: Sequence[Sequence[float]] | np.ndarray | None = None,
    classes: Sequence[Any] | np.ndarray | None = None,
    config: JointDistributionAdaptationConfig | dict[str, Any] | None = None,
    method: str | None = None,
    n_components: int | str | None = None,
    max_iterations: int | str | None = None,
    conditional_weight: float | str | None = None,
    regularization: float | str | None = None,
    eigen_ridge: float | str | None = None,
    temperature: float | str | None = None,
    standardize: bool | None = None,
    normalize_latent: bool | None = None,
) -> JointDistributionAdaptationResult:
    """Fit iterative marginal-plus-conditional source-target alignment.

    ``target_probabilities`` may contain source-model probabilities for the
    unlabeled target rows. If omitted, target pseudo-labels are initialized by
    nearest source-class centroids. No target-label argument exists.
    """

    cfg = _resolve_config(
        config,
        method=method,
        n_components=n_components,
        max_iterations=max_iterations,
        conditional_weight=conditional_weight,
        regularization=regularization,
        eigen_ridge=eigen_ridge,
        temperature=temperature,
        standardize=standardize,
        normalize_latent=normalize_latent,
    )
    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 the same feature width.")
    labels = _object_vector(source_labels, expected_length=source.shape[0], name="source_labels")
    class_order = _resolve_classes(labels, classes)
    if len(class_order) < 2:
        raise ValueError("Joint distribution adaptation requires at least two source classes.")
    encoded_source = _encode_labels(labels, class_order)

    joint = np.vstack([source, target]).astype(float, copy=False)
    mean = np.mean(joint, axis=0) if cfg.standardize else np.zeros(joint.shape[1], dtype=float)
    centered = joint - mean
    scale = np.std(centered, axis=0, ddof=1 if joint.shape[0] > 1 else 0) if cfg.standardize else np.ones(joint.shape[1], dtype=float)
    scale = np.maximum(scale, _MIN_SCALE)
    z = centered / scale
    source_z = z[: source.shape[0]]
    target_z = z[source.shape[0] :]

    if target_probabilities is None:
        responsibilities = _centroid_probabilities(source_z, encoded_source, target_z, len(class_order), temperature=cfg.temperature)
        used_initial_probabilities = False
    else:
        responsibilities = _probability_matrix(target_probabilities, expected_rows=target.shape[0], expected_classes=len(class_order))
        used_initial_probabilities = True
    pseudo = np.argmax(responsibilities, axis=1)

    centering = np.eye(joint.shape[0], dtype=float) - np.full((joint.shape[0], joint.shape[0]), 1.0 / float(joint.shape[0]))
    feature_dim = z.shape[1]
    component_count = _effective_components(cfg.n_components, n_samples=z.shape[0], n_features=feature_dim)
    converged = False
    projection = np.eye(feature_dim, component_count, dtype=float)
    selected_values = np.zeros(component_count, dtype=float)
    source_latent = source_z @ projection
    target_latent = target_z @ projection

    for iteration in range(1, cfg.max_iterations + 1):
        discrepancy = _marginal_matrix(source.shape[0], target.shape[0])
        discrepancy += cfg.conditional_weight * _conditional_matrix(encoded_source, responsibilities, len(class_order))
        norm = float(np.linalg.norm(discrepancy, ord="fro"))
        if norm > _MIN_SCALE:
            discrepancy /= norm
        a_matrix = z.T @ discrepancy @ z + cfg.regularization * np.eye(feature_dim, dtype=float)
        b_matrix = z.T @ centering @ z + cfg.eigen_ridge * np.eye(feature_dim, dtype=float)
        values, vectors = eigh(a_matrix, b_matrix, check_finite=True)
        selected = np.argsort(values)[:component_count]
        selected_values = values[selected]
        projection = _canonicalize_projection(vectors[:, selected])
        latent = z @ projection
        if cfg.normalize_latent:
            latent_scale = np.maximum(np.std(latent, axis=0, ddof=1 if latent.shape[0] > 1 else 0), _MIN_SCALE)
            latent /= latent_scale
            projection /= latent_scale.reshape(1, -1)
        source_latent = latent[: source.shape[0]]
        target_latent = latent[source.shape[0] :]
        updated_probabilities = _centroid_probabilities(
            source_latent,
            encoded_source,
            target_latent,
            len(class_order),
            temperature=cfg.temperature,
        )
        updated_pseudo = np.argmax(updated_probabilities, axis=1)
        converged = bool(np.array_equal(updated_pseudo, pseudo))
        pseudo = updated_pseudo
        responsibilities = updated_probabilities if cfg.method == "soft_jda" else _one_hot(pseudo, len(class_order))
        if converged:
            iterations_run = iteration
            break
    else:
        iterations_run = cfg.max_iterations

    pseudo_labels = _decode_labels(pseudo, class_order)
    metadata = _metadata(
        cfg=cfg,
        n_source_rows=source.shape[0],
        n_target_rows=target.shape[0],
        feature_dim=feature_dim,
        n_components=component_count,
        n_classes=len(class_order),
        iterations=iterations_run,
        converged=converged,
        used_initial_probabilities=used_initial_probabilities,
        pseudo=pseudo,
        eigenvalues=selected_values,
    )
    return JointDistributionAdaptationResult(
        source_features=source_latent.astype(np.float32, copy=False),
        target_features=target_latent.astype(np.float32, copy=False),
        projection=projection.astype(np.float32, copy=False),
        feature_mean=mean.astype(np.float32, copy=False),
        feature_scale=scale.astype(np.float32, copy=False),
        eigenvalues=np.asarray(selected_values, dtype=float),
        target_pseudo_labels=pseudo_labels,
        target_probabilities=np.asarray(responsibilities, dtype=np.float32),
        classes=class_order,
        n_iterations=int(iterations_run),
        converged=converged,
        metadata=metadata,
    )

transform_joint_distribution_features(features, result)

Transform new rows with a fitted JDA projection.

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

def transform_joint_distribution_features(features: Sequence[Sequence[float]] | np.ndarray, result: JointDistributionAdaptationResult) -> np.ndarray:
    """Transform new rows with a fitted JDA projection."""

    matrix = _feature_matrix(features, name="features")
    if matrix.shape[1] != result.projection.shape[0]:
        raise ValueError("features width does not match the fitted projection.")
    return (((matrix - result.feature_mean) / result.feature_scale) @ result.projection).astype(np.float32, copy=False)

joint_distribution_adaptation_config(*, method='jda', n_components=16, max_iterations=10, conditional_weight=1.0, regularization=0.001, eigen_ridge=1e-06, temperature=1.0, standardize=True, normalize_latent=False)

Normalize JDA configuration values.

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

def joint_distribution_adaptation_config(
    *,
    method: str | None = "jda",
    n_components: int | str | None = 16,
    max_iterations: int | str = 10,
    conditional_weight: float | str = 1.0,
    regularization: float | str = 1e-3,
    eigen_ridge: float | str = 1e-6,
    temperature: float | str = 1.0,
    standardize: bool = True,
    normalize_latent: bool = False,
) -> JointDistributionAdaptationConfig:
    """Normalize JDA configuration values."""

    return JointDistributionAdaptationConfig(
        method=normalize_jda_method(method),
        n_components=_normalize_components(n_components),
        max_iterations=_positive_int(max_iterations, name="max_iterations"),
        conditional_weight=_nonnegative_float(conditional_weight, name="conditional_weight"),
        regularization=_nonnegative_float(regularization, name="regularization"),
        eigen_ridge=_positive_float(eigen_ridge, name="eigen_ridge"),
        temperature=_positive_float(temperature, name="temperature"),
        standardize=bool(standardize),
        normalize_latent=bool(normalize_latent),
    )

normalize_jda_method(method)

Normalize public JDA method aliases.

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

def normalize_jda_method(method: str | None) -> str:
    """Normalize public JDA method aliases."""

    normalized = "jda" if method is None else str(method).strip().lower().replace("-", "_")
    normalized = {
        "joint_distribution_adaptation": "jda",
        "hard_jda": "jda",
        "soft": "soft_jda",
        "probabilistic_jda": "soft_jda",
        "soft_joint_distribution_adaptation": "soft_jda",
    }.get(normalized, normalized)
    if normalized not in JDA_METHODS:
        raise ValueError(f"Unknown JDA method {method!r}. Available methods: {', '.join(JDA_METHODS)}.")
    return normalized