Skip to content

Unlabeled prior-shift adaptation

neureptrace.decoding.prior_shift adapts target probability traces when a source-trained decoder is evaluated on a held-out target batch with a different class prior.

The method is Category 2 / unlabeled target-adaptive. It uses target probability rows, and optionally a source training prior, but it does not accept held-out target labels.

Typical use:

from neureptrace.decoding.prior_shift import adapt_probabilities_for_prior_shift

result = adapt_probabilities_for_prior_shift(
    target_probabilities,
    source_prior=[0.5, 0.5],
)
adapted_probabilities = result.probabilities

For run-wise or block-wise shifts, use adapt_probability_blocks_for_prior_shift with a block id per target row. The block-wise protocol is still target-label-free, but it should be reported separately from strict source-only decoding because the held-out target probability distribution is used for adaptation.

neureptrace.decoding.prior_shift

Unlabeled target prior-shift adaptation for probability traces.

PriorShiftAdaptationResult dataclass

Probability rows after unlabeled prior-shift adaptation.

Source code in src/neureptrace/decoding/prior_shift.py
16
17
18
19
20
21
22
23
24
25
26
27
@dataclass(frozen=True, slots=True)
class PriorShiftAdaptationResult:
    """Probability rows after unlabeled prior-shift adaptation."""

    probabilities: np.ndarray
    target_prior: np.ndarray
    source_prior: np.ndarray
    class_bias: np.ndarray
    n_iterations: int
    converged: bool
    max_delta: float
    metadata: dict[str, Any] = field(default_factory=dict)

PriorShiftBlockResult dataclass

Block-wise prior-shift adaptation output.

Source code in src/neureptrace/decoding/prior_shift.py
30
31
32
33
34
35
36
@dataclass(frozen=True, slots=True)
class PriorShiftBlockResult:
    """Block-wise prior-shift adaptation output."""

    probabilities: np.ndarray
    block_results: Mapping[Hashable, PriorShiftAdaptationResult]
    metadata: dict[str, Any] = field(default_factory=dict)

adapt_probabilities_for_prior_shift(probabilities, *, source_prior=None, initial_target_prior=None, target_prior=None, max_iter=100, tol=1e-08, smoothing=1e-06, damping=1.0, epsilon=EPSILON)

Estimate a target class prior from unlabeled probability rows and reweight.

The function accepts source-model posteriors for target rows and optionally a source training prior. It never accepts target class labels. If target_prior is supplied, EM is skipped and only prior-ratio reweighting is applied.

Source code in src/neureptrace/decoding/prior_shift.py
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
def adapt_probabilities_for_prior_shift(
    probabilities: Sequence[Sequence[float]] | np.ndarray,
    *,
    source_prior: Sequence[float] | np.ndarray | None = None,
    initial_target_prior: Sequence[float] | np.ndarray | None = None,
    target_prior: Sequence[float] | np.ndarray | None = None,
    max_iter: int | str = 100,
    tol: float | str = 1e-8,
    smoothing: float | str = 1e-6,
    damping: float | str = 1.0,
    epsilon: float | str = EPSILON,
) -> PriorShiftAdaptationResult:
    """Estimate a target class prior from unlabeled probability rows and reweight.

    The function accepts source-model posteriors for target rows and optionally a
    source training prior. It never accepts target class labels. If ``target_prior``
    is supplied, EM is skipped and only prior-ratio reweighting is applied.
    """

    matrix = _probability_matrix(probabilities, epsilon=epsilon)
    n_rows, n_classes = matrix.shape
    eps = _positive_float(epsilon, name="epsilon")
    source = _prior(source_prior, n_classes=n_classes, default="uniform", name="source_prior", epsilon=eps)
    smooth = _nonnegative_float(smoothing, name="smoothing")
    max_iterations = _positive_int(max_iter, name="max_iter")
    tolerance = _nonnegative_float(tol, name="tol")
    damp = _unit_interval_float(damping, name="damping")

    if target_prior is not None:
        target = _prior(target_prior, n_classes=n_classes, default="uniform", name="target_prior", epsilon=eps)
        iterations = 0
        converged = True
        max_delta = 0.0
        mode = "fixed_target_prior"
    else:
        if initial_target_prior is None:
            target = _normalize_prior(np.mean(matrix, axis=0), epsilon=eps)
        else:
            target = _prior(initial_target_prior, n_classes=n_classes, default="uniform", name="initial_target_prior", epsilon=eps)
        target = _smooth(target, smooth, eps)
        iterations = 0
        converged = False
        max_delta = float("inf")
        mode = "em_estimated_target_prior"
        for iterations in range(1, max_iterations + 1):
            responsibilities = reweight_probabilities_by_prior(matrix, source_prior=source, target_prior=target, epsilon=eps)
            updated = _smooth(np.mean(responsibilities, axis=0), smooth, eps)
            if damp < 1.0:
                updated = _normalize_prior((1.0 - damp) * target + damp * updated, epsilon=eps)
            max_delta = float(np.max(np.abs(updated - target)))
            target = updated
            if max_delta <= tolerance:
                converged = True
                break

    adapted = reweight_probabilities_by_prior(matrix, source_prior=source, target_prior=target, epsilon=eps)
    class_bias = _normalize_bias(target / np.maximum(source, eps))
    metadata = _metadata(
        mode=mode,
        n_rows=n_rows,
        n_classes=n_classes,
        source_prior=source,
        target_prior=target,
        class_bias=class_bias,
        iterations=iterations,
        converged=converged,
        max_delta=max_delta,
        blockwise=False,
    )
    return PriorShiftAdaptationResult(adapted, target, source, class_bias, int(iterations), bool(converged), float(max_delta), metadata)

adapt_probability_blocks_for_prior_shift(probabilities, block_ids, *, source_prior=None, min_block_rows=2, **kwargs)

Run prior-shift adaptation separately for each target block.

Source code in src/neureptrace/decoding/prior_shift.py
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
def adapt_probability_blocks_for_prior_shift(
    probabilities: Sequence[Sequence[float]] | np.ndarray,
    block_ids: Sequence[Hashable] | np.ndarray,
    *,
    source_prior: Sequence[float] | np.ndarray | None = None,
    min_block_rows: int | str = 2,
    **kwargs: Any,
) -> PriorShiftBlockResult:
    """Run prior-shift adaptation separately for each target block."""

    matrix = _probability_matrix(probabilities, epsilon=kwargs.get("epsilon", EPSILON))
    blocks = _object_vector(block_ids, expected_length=matrix.shape[0], name="block_ids")
    minimum = _positive_int(min_block_rows, name="min_block_rows")
    source_prior = _materialize_numeric_iterables(source_prior)
    kwargs = _materialize_reused_prior_kwargs(kwargs)
    adapted = np.empty_like(matrix)
    results: dict[Hashable, PriorShiftAdaptationResult] = {}
    for block in _unique_values(blocks):
        mask = _object_equal_mask(blocks, block)
        if int(np.sum(mask)) < minimum:
            raise ValueError(f"Block {block!r} has fewer than min_block_rows={minimum} rows.")
        result = adapt_probabilities_for_prior_shift(matrix[mask], source_prior=source_prior, **kwargs)
        adapted[mask] = result.probabilities
        results[block] = result
    metadata = {
        "prior_shift_adaptation": True,
        "prior_shift_protocol": PRIOR_SHIFT_PROTOCOL,
        "prior_shift_protocol_category": PRIOR_SHIFT_CATEGORY,
        "prior_shift_uses_target_probabilities": True,
        "prior_shift_uses_target_labels": False,
        "prior_shift_blockwise": True,
        "prior_shift_n_rows": int(matrix.shape[0]),
        "prior_shift_n_classes": int(matrix.shape[1]),
        "prior_shift_n_blocks": int(len(results)),
        "prior_shift_blocks": "|".join(str(block) for block in results),
        "prior_shift_converged_all_blocks": all(result.converged for result in results.values()),
        "prior_shift_target_priors_by_block": "|".join(f"{block}:{_format(result.target_prior)}" for block, result in results.items()),
    }
    return PriorShiftBlockResult(adapted, results, metadata)

reweight_probabilities_by_prior(probabilities, *, source_prior, target_prior, epsilon=EPSILON)

Reweight posterior rows by target_prior / source_prior.

Source code in src/neureptrace/decoding/prior_shift.py
111
112
113
114
115
116
117
118
119
120
121
122
123
124
def reweight_probabilities_by_prior(
    probabilities: Sequence[Sequence[float]] | np.ndarray,
    *,
    source_prior: Sequence[float] | np.ndarray,
    target_prior: Sequence[float] | np.ndarray,
    epsilon: float | str = EPSILON,
) -> np.ndarray:
    """Reweight posterior rows by ``target_prior / source_prior``."""

    matrix = _probability_matrix(probabilities, epsilon=epsilon)
    eps = _positive_float(epsilon, name="epsilon")
    source = _prior(source_prior, n_classes=matrix.shape[1], default="uniform", name="source_prior", epsilon=eps)
    target = _prior(target_prior, n_classes=matrix.shape[1], default="uniform", name="target_prior", epsilon=eps)
    return _normalize_rows(matrix * (target / np.maximum(source, eps))[None, :], epsilon=eps)

prior_from_labels(labels, classes=None, *, smoothing=0.0)

Compute an empirical prior from source labels.

Source code in src/neureptrace/decoding/prior_shift.py
168
169
170
171
172
173
174
175
176
def prior_from_labels(labels: Sequence[Hashable] | np.ndarray, classes: Sequence[Hashable] | np.ndarray | None = None, *, smoothing: float | str = 0.0) -> tuple[np.ndarray, tuple[Hashable, ...]]:
    """Compute an empirical prior from source labels."""

    label_vector = _object_vector(labels, name="labels")
    if label_vector.size == 0:
        raise ValueError("labels must contain at least one row.")
    class_order = _unique_values(label_vector) if classes is None else tuple(_object_vector(classes, name="classes").tolist())
    counts = np.asarray([np.count_nonzero(_object_equal_mask(label_vector, class_label)) for class_label in class_order], dtype=float)
    return _normalize_prior(counts + _nonnegative_float(smoothing, name="smoothing"), epsilon=EPSILON), class_order