| """Lecture 7: every number the chapter works by hand, asserted against the file that prints it. |
| |
| Chapter 7 computes almost all of its results on paper, so these tests are the chapter itself read |
| back. The tolerances are the number of digits the chapter quotes. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import importlib.util |
| import math |
| import sys |
| from pathlib import Path |
|
|
| import pytest |
| import torch |
|
|
| LECTURE = Path(__file__).resolve().parent.parent / "lecture_7" |
|
|
|
|
| def load(name: str): |
| """Import one lecture file by path, since ``lecture_7`` is a folder of scripts, not a package.""" |
| spec = importlib.util.spec_from_file_location(f"lecture_7_{name}", LECTURE / f"{name}.py") |
| module = importlib.util.module_from_spec(spec) |
| sys.modules[spec.name] = module |
| spec.loader.exec_module(module) |
| return module |
|
|
|
|
| def tensor(*values): |
| return torch.tensor(values, dtype=torch.float64) |
|
|
|
|
| |
|
|
|
|
| def test_four_points_and_two_pairings(): |
| ot = load("optimal_transport") |
| pairings = ot.two_pairings() |
| assert pairings["arrows_ordered"] == pytest.approx(2.0) |
| assert pairings["arrows_crossed"] == pytest.approx(10.0) |
| assert pairings["mass_ordered"] == pytest.approx(1.0) |
| assert pairings["mass_crossed"] == pytest.approx(5.0) |
| assert pairings["mass_independent"] == pytest.approx(3.0) |
| assert pairings["C"].tolist() == [[1.0, 9.0], [1.0, 1.0]] |
|
|
|
|
| def test_two_by_two_plan_and_prices(): |
| ot = load("optimal_transport") |
| assert ot.feasible_interval() == (0.0, 0.3) |
| C = ot.cost_matrix(ot.SOURCE_X, ot.TARGET_Y) |
| assert ot.cost_along_the_segment(0.0, C) == pytest.approx(5.8) |
| assert ot.cost_along_the_segment(0.3, C) == pytest.approx(3.4) |
| plan = ot.feasible_plan(0.3) |
| assert plan.tolist() == [[0.3, 0.3], [0.0, 0.4]] |
| assert ot.is_coupling(plan, ot.SOURCE_A, ot.TARGET_B) |
| f, g = ot.prices_from_loaded_routes(C, plan) |
| assert f.tolist() == [0.0, -8.0] and g.tolist() == [1.0, 9.0] |
| assert ot.slack(C, f, g).tolist() == [[0.0, 0.0], [8.0, 0.0]] |
| assert ot.dual_value(ot.SOURCE_A, f, ot.TARGET_B, g) == pytest.approx(3.4) |
| assert ot.plan_cost(C, plan) ** 0.5 == pytest.approx(1.844, abs=5e-4) |
|
|
|
|
| def test_no_crossing_and_three_points(): |
| ot = load("optimal_transport") |
| expanded, collected = ot.crossing_difference(0.0, 2.0, 1.0, 3.0) |
| assert expanded == pytest.approx(collected) and collected == pytest.approx(-8.0) |
| x = tensor(4.0, 1.0, 6.0) |
| y = tensor(9.0, 2.0, 3.0) |
| costs = [c for c, _ in ot.monge_permutations(x.sort().values, y.sort().values)] |
| assert costs == pytest.approx([11 / 3, 17 / 3, 35 / 3, 45 / 3, 77 / 3, 81 / 3]) |
| assert ot.wasserstein_on_the_line(x, y, 2.0, 3072) == pytest.approx(11 / 3) |
|
|
|
|
| def test_gaussian_closed_form(): |
| ot = load("optimal_transport") |
| assert ot.gaussian_w2_squared(0.0, 1.0, 3.0, 2.0) == pytest.approx(10.0) |
| assert math.sqrt(10.0) == pytest.approx(3.162, abs=5e-4) |
| points = tensor(-1.0, 0.0, 1.0) |
| assert ot.brenier_map(points, 0.0, 1.0, 3.0, 2.0).tolist() == [1.0, 3.0, 5.0] |
| source = torch.special.ndtri(ot.quantile_levels(3072)) |
| mapped = ot.brenier_map(source, 0.0, 1.0, 3.0, 2.0) |
| assert float(((mapped - source) ** 2).mean()) == pytest.approx(10.0, abs=1e-3) |
|
|
|
|
| |
|
|
|
|
| def test_entropies_of_the_three_couplings(): |
| sk = load("sinkhorn") |
| independent = sk.independent_coupling(sk.SOURCE_A, sk.TARGET_B) |
| assert independent.flatten().tolist() == pytest.approx([0.18, 0.42, 0.12, 0.28]) |
| assert sk.entropy(sk.EXACT_PLAN) == pytest.approx(1.089, abs=5e-4) |
| assert sk.entropy(independent) == pytest.approx(1.284, abs=5e-4) |
| assert sk.entropy(sk.SOURCE_A) == pytest.approx(0.673, abs=5e-4) |
| assert sk.entropy(sk.TARGET_B) == pytest.approx(0.611, abs=5e-4) |
| assert sk.relative_entropy(sk.EXACT_PLAN, independent) == pytest.approx(0.195, abs=5e-4) |
|
|
|
|
| def test_sinkhorn_first_sweep_by_hand(): |
| sk = load("sinkhorn") |
| K = sk.gibbs_kernel(sk.COST, 2.0) |
| assert K.flatten().tolist() == pytest.approx([0.606531, 0.011109, 0.606531, 0.606531], |
| abs=5e-7) |
| v = torch.ones(2, dtype=torch.float64) |
| assert (K @ v).tolist() == pytest.approx([0.617640, 1.213061], abs=5e-7) |
| u = sk.SOURCE_A / (K @ v) |
| assert u.tolist() == pytest.approx([0.971440, 0.329744], abs=5e-7) |
| assert (K.T @ u).tolist() == pytest.approx([0.789208, 0.210792], abs=5e-7) |
| v = sk.TARGET_B / (K.T @ u) |
| assert v.tolist() == pytest.approx([0.38013, 3.32081], abs=5e-6) |
| plan = sk.scaling_form(u, K, v) |
| assert plan.flatten().tolist() == pytest.approx([0.2240, 0.0358, 0.0760, 0.6642], abs=5e-5) |
| assert plan.sum(dim=0).tolist() == pytest.approx([0.3, 0.7]) |
| assert plan.sum(dim=1).tolist() == pytest.approx([0.2598, 0.7402], abs=5e-5) |
| after = sk.project_rows(plan, sk.SOURCE_A) |
| assert after.sum(dim=0).tolist() == pytest.approx([0.558, 0.442], abs=5e-4) |
|
|
|
|
| def test_sinkhorn_convergence_and_epsilon(): |
| sk = load("sinkhorn") |
| K = sk.gibbs_kernel(sk.COST, 2.0) |
| u, v, errors = sk.sinkhorn_sweeps(K, sk.SOURCE_A, sk.TARGET_B, 40) |
| assert errors[:5] == pytest.approx([0.680, 0.382, 0.155, 0.047, 0.012], abs=5e-4) |
| assert sk.sweeps_to_tolerance(K, sk.SOURCE_A, sk.TARGET_B) == 17 |
| plan = sk.scaling_form(u, K, v) |
| assert plan.flatten().tolist() == pytest.approx( |
| [0.293122, 0.306878, 0.006878, 0.393122], abs=5e-7) |
| assert float((sk.COST * plan).sum()) == pytest.approx(3.45502, abs=5e-6) |
| assert sk.entropy(plan) == pytest.approx(1.124, abs=5e-4) |
| independent = sk.independent_coupling(sk.SOURCE_A, sk.TARGET_B) |
| assert sk.entropic_objective(plan, sk.COST, 2.0) == pytest.approx(1.208, abs=5e-4) |
| assert sk.entropic_objective(sk.EXACT_PLAN, sk.COST, 2.0) == pytest.approx(1.222, abs=5e-4) |
| assert sk.entropic_objective(independent, sk.COST, 2.0) == pytest.approx(1.792, abs=5e-4) |
| assert sk.relative_entropy(plan, independent) == pytest.approx(0.160, abs=5e-4) |
| f = sk.multipliers_from_scalings(u, 2.0) |
| g = sk.multipliers_from_scalings(v, 2.0) |
| assert f.tolist() == pytest.approx([3.766, -3.739], abs=5e-4) |
| assert g.tolist() == pytest.approx([-3.220, 4.872], abs=5e-4) |
| expected = {10.0: (4.0540165, 7), 2.0: (3.4550204, 17), 0.5: (3.4000004, 28), |
| 0.1: (3.4000000, 80)} |
| for epsilon, (cost, sweeps) in expected.items(): |
| kernel = sk.gibbs_kernel(sk.COST, epsilon) |
| su, sv, _ = sk.sinkhorn_sweeps(kernel, sk.SOURCE_A, sk.TARGET_B, 400) |
| assert float((sk.COST * sk.scaling_form(su, kernel, sv)).sum()) == pytest.approx( |
| cost, abs=5e-8) |
| assert sk.sweeps_to_tolerance(kernel, sk.SOURCE_A, sk.TARGET_B) == sweeps |
|
|
|
|
| |
|
|
|
|
| def test_one_particle_and_the_averaged_bound(): |
| dp = load("displacement") |
| times = torch.linspace(0.0, 1.0, 1025, dtype=torch.float64) |
| assert dp.action_of_a_path(dp.steady_particle(times), times) == pytest.approx(4.0) |
| assert dp.action_of_a_path(dp.dawdling_particle(times), times) == pytest.approx(8.0) |
| x, y = tensor(0.0, 2.0), tensor(1.0, 3.0) |
| ordered = torch.tensor([[0.5, 0.0], [0.0, 0.5]], dtype=torch.float64) |
| crossed = torch.tensor([[0.0, 0.5], [0.5, 0.0]], dtype=torch.float64) |
| independent = torch.full((2, 2), 0.25, dtype=torch.float64) |
| assert dp.average_squared_displacement(x, y, ordered) == pytest.approx(1.0) |
| assert dp.average_squared_displacement(x, y, independent) == pytest.approx(3.0) |
| assert dp.average_squared_displacement(x, y, crossed) == pytest.approx(5.0) |
|
|
|
|
| def test_displacement_interpolation_and_its_velocity(): |
| dp = load("displacement") |
| mean, std = dp.displacement_law(0.5) |
| assert (mean, std) == pytest.approx((1.5, 1.5)) |
| assert dp.gaussian_w2(0.0, 1.0, mean, std) == pytest.approx(math.sqrt(2.5)) |
| assert math.sqrt(2.5) == pytest.approx(1.581, abs=5e-4) |
| assert dp.gaussian_w2(0.0, 1.0, 3.0, 2.0) == pytest.approx(3.162, abs=5e-4) |
| grid = torch.linspace(-6.0, 9.0, 15001, dtype=torch.float64) |
| peak = float(grid[int(dp.mixture_density(grid).argmax())]) |
| assert peak == pytest.approx(0.13, abs=5e-3) |
| half = torch.tensor(0.5, dtype=torch.float64) |
| assert float(dp.velocity_field(torch.tensor(1.5, dtype=torch.float64), half)) == pytest.approx(3.0) |
| assert float(dp.velocity_field(torch.tensor(3.0, dtype=torch.float64), half)) == pytest.approx(4.0) |
| for t in (0.0, 0.25, 0.5, 0.75, 1.0): |
| assert dp.action_density(t, 4096) == pytest.approx(10.0, abs=1e-3) |
| assert abs(dp.continuity_residual(1.5, 0.5)) < 1e-12 |
|
|
|
|
| |
|
|
|
|
| def test_batch_coupling_costs(): |
| fm = load("ot_flow_matching") |
| x = torch.tensor([[0.0], [2.0]]) |
| y = torch.tensor([[1.0], [3.0]]) |
| C = fm.batch_cost_matrix(x, y) |
| assert C.flatten().tolist() == pytest.approx([1.0, 9.0, 1.0, 1.0]) |
| assert float(C.mean()) == pytest.approx(3.0) |
| pairing = fm.optimal_pairing(x, y) |
| assert pairing.tolist() == [0, 1] |
| assert float(C[torch.arange(2), pairing].mean()) == pytest.approx(1.0) |
| assert fm.vertex_count(256, 256) == {"unknowns": 65536, "constraints": 512, |
| "loaded_routes": 511} |
| assert fm.conditional_velocity_energy(x, fm.couple(x, y, "optimal")) == pytest.approx(1.0) |
|
|
|
|
| def test_flow_matching_pipeline_runs(): |
| fm = load("ot_flow_matching") |
| from cis6270.nets import VelocityMLP |
|
|
| torch.manual_seed(0) |
| model = VelocityMLP(dim=2, width=16) |
| x0 = torch.randn(8, 2) |
| x1 = torch.randn(8, 2) |
| assert float(fm.flow_matching_loss(model, x0, fm.couple(x0, x1, "optimal"))) > 0.0 |
| samples, straightness = fm.integrate(model, x0, 5) |
| assert samples.shape == (8, 2) and straightness >= 0.0 |
|
|
|
|
| |
|
|
|
|
| def test_eight_paths_and_the_chain_rule(): |
| sb = load("static_bridge") |
| law = sb.path_law(sb.PI_0, [sb.A, sb.A]) |
| assert law[(0, 0, 1)] == pytest.approx(0.105) |
| assert law[(0, 0, 0)] == pytest.approx(0.245) |
| assert law[(1, 1, 1)] == pytest.approx(0.180) |
| assert sum(law.values()) == pytest.approx(1.0) |
| table = torch.tensor([[0.5, 0.1], [0.1, 0.3]], dtype=torch.float64) |
| uniform = torch.full((2, 2), 0.25, dtype=torch.float64) |
| total, marginal, conditional = sb.chain_rule_terms(table, uniform) |
| assert total == pytest.approx(0.2180, abs=5e-5) |
| assert marginal == pytest.approx(0.0201, abs=5e-5) |
| assert conditional == pytest.approx(0.1979, abs=5e-5) |
| assert marginal + conditional == pytest.approx(total) |
|
|
|
|
| def test_proportional_fitting_on_the_chain(): |
| sb = load("static_bridge") |
| R01 = sb.reference_endpoint_law(sb.PI_0, sb.A, 2) |
| assert R01.flatten().tolist() == pytest.approx([0.305, 0.195, 0.26, 0.24]) |
| assert float((R01.sum(dim=0) - sb.PI_1).abs().sum()) == pytest.approx(0.73) |
| f, g, errors = sb.proportional_fitting(R01, sb.PI_0, sb.PI_1, 4) |
| assert errors == pytest.approx([3.933e-3, 2.098e-5, 1.119e-7, 5.966e-10], rel=1e-3) |
| assert f.tolist() == pytest.approx([1.072002, 0.937019], abs=5e-7) |
| assert g.tolist() == pytest.approx([0.350517, 1.843637], abs=5e-7) |
| plan = f[:, None] * R01 * g[None, :] |
| assert plan.flatten().tolist() == pytest.approx( |
| [0.114605, 0.385395, 0.085395, 0.414605], abs=5e-7) |
| assert sb.kl_vectors(plan.flatten(), R01.flatten()) == pytest.approx(0.2820, abs=5e-5) |
| independent = sb.PI_0[:, None] * sb.PI_1[None, :] |
| assert sb.kl_vectors(independent.flatten(), R01.flatten()) == pytest.approx(0.2847, abs=5e-5) |
|
|
|
|
| def test_h_transform_and_the_bridge_path_law(): |
| sb = load("static_bridge") |
| R01 = sb.reference_endpoint_law(sb.PI_0, sb.A, 2) |
| f, g, _ = sb.proportional_fitting(R01, sb.PI_0, sb.PI_1, 40) |
| potentials = sb.backward_potential(sb.A, g, 2) |
| assert potentials[2].tolist() == pytest.approx([0.350517, 1.843637], abs=5e-7) |
| assert potentials[1].tolist() == pytest.approx([0.798453, 1.246389], abs=5e-7) |
| assert potentials[0].tolist() == pytest.approx([0.932834, 1.067215], abs=5e-7) |
| tilted = [sb.tilted_matrix(sb.A, potentials[k], potentials[k + 1]) for k in range(2)] |
| assert tilted[0].flatten().tolist() == pytest.approx( |
| [0.599160, 0.400840, 0.299266, 0.700734], abs=5e-7) |
| assert tilted[1].flatten().tolist() == pytest.approx( |
| [0.307297, 0.692703, 0.112490, 0.887510], abs=5e-7) |
| assert tilted[0].sum(dim=1).tolist() == pytest.approx([1.0, 1.0]) |
| middle = sb.PI_0 @ tilted[0] |
| assert middle.tolist() == pytest.approx([0.449213, 0.550787], abs=5e-7) |
| assert (middle @ tilted[1]).tolist() == pytest.approx([0.2, 0.8]) |
| assert (1.0 / potentials[0]).tolist() == pytest.approx(f.tolist(), abs=1e-9) |
| plan = f[:, None] * R01 * g[None, :] |
| bridge = sb.fill_with_reference_bridges(plan, sb.A) |
| reference = sb.path_law(sb.PI_0, [sb.A, sb.A]) |
| assert bridge[(0, 0, 0)] == pytest.approx(0.092060, abs=5e-7) |
| assert bridge[(0, 0, 1)] == pytest.approx(0.207520, abs=5e-7) |
| assert bridge[(1, 1, 1)] == pytest.approx(0.310954, abs=5e-7) |
| ratios = {round(bridge[p] / reference[p], 6) for p in reference} |
| assert ratios == {0.375755, 1.976383, 0.328441, 1.727522} |
| assert sb.relative_entropy(bridge, reference) == pytest.approx(0.281961, abs=5e-7) |
| assert sum(v for k, v in bridge.items() if k[1] == 0) == pytest.approx(0.449213, abs=5e-7) |
|
|
|
|
| def test_brownian_reference_recovers_the_entropic_plan(): |
| sb = load("static_bridge") |
| plan = sb.brownian_reference_plan(1.0) |
| assert plan.flatten().tolist() == pytest.approx([0.2931, 0.3069, 0.0069, 0.3931], abs=5e-5) |
|
|
|
|
| |
|
|
|
|
| def test_constant_control_two_ways(): |
| db = load("dynamic_bridge") |
| assert db.control_energy_entropy(2.0, 1.0) == pytest.approx(2.0) |
| assert db.gaussian_entropy_gap(2.0, 1.0) == pytest.approx(2.0) |
| generator = torch.Generator().manual_seed(0) |
| _, log_ratio = db.simulate_controlled(2.0, 200, 8192, 1.0, generator) |
| assert float(log_ratio.mean()) == pytest.approx(2.0, abs=0.1) |
|
|
|
|
| def test_the_bridge_potential_solves_both_equations(): |
| db = load("dynamic_bridge") |
| value = db.hopf_cole(db.bridge_potential) |
| for x, t in ((-0.5, 0.2), (0.5, 0.5), (2.5, 0.8)): |
| assert abs(db.backward_residual(db.bridge_potential, x, t)) < 1e-12 |
| assert abs(db.hjb_residual(value, x, t)) < 1e-12 |
| drift = float(db.bridge_drift(torch.tensor(x, dtype=torch.float64), |
| torch.tensor(t, dtype=torch.float64))) |
| assert db.optimal_control(value, x, t) == pytest.approx(drift, abs=1e-10) |
|
|
|
|
| def test_brownian_bridge_numbers(): |
| db = load("dynamic_bridge") |
| mean, variance = db.bridge_law(0.25) |
| assert mean == pytest.approx(0.5) |
| assert variance == pytest.approx(0.1875) |
| assert math.sqrt(variance) == pytest.approx(0.4330, abs=5e-5) |
| drift = float(db.bridge_drift(torch.tensor(0.6732, dtype=torch.float64), |
| torch.tensor(0.25, dtype=torch.float64))) |
| assert drift == pytest.approx(1.7691, abs=5e-5) |
| assert db.bridge_law(0.5)[1] == pytest.approx(0.25) |
| for epsilon, spread in ((0.25, 0.25), (1.0, 0.5), (4.0, 1.0)): |
| assert math.sqrt(db.bridge_law(0.5, epsilon=epsilon)[1]) == pytest.approx(spread) |
|
|
|
|
| |
|
|
|
|
| def test_dsb_half_steps_on_the_chain(): |
| dsb = load("dsb") |
| reference = dsb.reference_path_law(dsb.PI_0, dsb.A, 2) |
| assert dsb.marginal(reference, 1).tolist() == pytest.approx([0.55, 0.45]) |
| assert dsb.marginal(reference, 2).tolist() == pytest.approx([0.565, 0.435]) |
| B2 = dsb.reverse_step(reference, 1) |
| assert B2.flatten().tolist() == pytest.approx( |
| [0.681416, 0.318584, 0.379310, 0.620690], abs=5e-7) |
| assert float(dsb.reverse_step(reference, 0)[0, 0]) == pytest.approx(0.636364, abs=5e-7) |
| half = dsb.fix_terminal(reference, dsb.PI_1) |
| assert dsb.marginal(half, 2).tolist() == pytest.approx([0.2, 0.8]) |
| assert dsb.marginal(half, 1).tolist() == pytest.approx([0.439731, 0.560269], abs=5e-7) |
| assert dsb.marginal(half, 0).tolist() == pytest.approx([0.466585, 0.533415], abs=5e-7) |
| assert float((dsb.marginal(half, 0) - dsb.PI_0).abs().sum()) == pytest.approx(0.0668, abs=5e-5) |
| full = dsb.fix_initial(half, dsb.PI_0) |
| assert dsb.one_step_conditional(full, 0).flatten().tolist() == pytest.approx( |
| [0.599738, 0.400262, 0.299771, 0.700229], abs=5e-7) |
| assert dsb.one_step_conditional(full, 1).flatten().tolist() == pytest.approx( |
| [0.309924, 0.690076, 0.113725, 0.886275], abs=5e-7) |
| assert dsb.marginal(full, 2).tolist() == pytest.approx([0.201967, 0.798033], abs=5e-7) |
| gap = max(float((dsb.one_step_conditional(full, 0) - dsb.TILTED_0).abs().max()), |
| float((dsb.one_step_conditional(full, 1) - dsb.TILTED_1).abs().max())) |
| assert gap < 0.003 |
| _, errors, _ = dsb.proportional_fitting(reference, dsb.PI_0, dsb.PI_1, 4) |
| assert errors == pytest.approx([3.933e-3, 2.098e-5, 1.119e-7, 5.966e-10], rel=1e-3) |
| assert errors[0] / errors[1] == pytest.approx(187.5, abs=0.1) |
|
|
|
|
| |
|
|
|
|
| def test_markov_drift_at_one_point(): |
| dm = load("dsbm") |
| assert dm.bridge_drift(3.0, 1.0, 0.5) == pytest.approx(4.0) |
| assert dm.bridge_drift(0.0, 1.0, 0.5) == pytest.approx(-2.0) |
| endpoints = tensor(3.0, 0.0) |
| weights = tensor(0.5, 0.5) |
| assert dm.markov_drift(endpoints, weights, 1.0, 0.5) == pytest.approx(1.0) |
| assert dm.least_squares_derivative(1.0, endpoints, weights, 1.0, 0.5) == pytest.approx(0.0) |
| assert dm.least_squares_derivative(0.0, endpoints, weights, 1.0, 0.5) == pytest.approx(-2.0) |
|
|
|
|
| def test_one_imf_round_on_the_chain(): |
| dm = load("dsbm") |
| bridges = dm.reference_bridge_table(dm.A) |
| assert float(bridges[0, 0, 0]) == pytest.approx(0.803279, abs=5e-7) |
| assert float(bridges[0, 0, 1]) == pytest.approx(0.538462, abs=5e-7) |
| target = dm.bridge_coupling(dm.A, dm.PI_0, dm.PI_1) |
| independent = dm.PI_0[:, None] * dm.PI_1[None, :] |
| assert independent.flatten().tolist() == pytest.approx([0.1, 0.4, 0.1, 0.4]) |
| assert float((independent - target).abs().sum()) == pytest.approx(0.0584, abs=5e-5) |
| coupling, gaps, first = dm.imf_rounds(independent, dm.A, 9, target) |
| assert first[0].flatten().tolist() == pytest.approx( |
| [0.591425, 0.408575, 0.307692, 0.692308], abs=5e-7) |
| assert first[1].flatten().tolist() == pytest.approx( |
| [0.298457, 0.701543, 0.119588, 0.880412], abs=5e-7) |
| projected = torch.diag(dm.PI_0) @ first[0] @ first[1] |
| assert projected.flatten().tolist() == pytest.approx( |
| [0.112688, 0.387312, 0.087312, 0.412688], abs=5e-7) |
| assert projected.sum(dim=0).tolist() == pytest.approx([0.2, 0.8]) |
| assert projected.sum(dim=1).tolist() == pytest.approx([0.5, 0.5]) |
| assert gaps[0] == pytest.approx(0.0077, abs=5e-5) |
| assert gaps[3] / gaps[4] == pytest.approx(7.37, abs=5e-3) |
| assert gaps[8] < 1e-9 |
|
|
|
|
| |
|
|
|
|
| def test_the_two_regression_targets(): |
| sm = load("sf2m") |
| x0, x1 = tensor(0.0)[0], tensor(2.0)[0] |
| z = tensor(0.4)[0] |
| for t_value, mean, sample, drift, flow, coefficient in ( |
| (0.25, 0.5, 0.6732, 1.7691, 2.2309, 4 / 3), |
| (0.75, 1.5, 1.6732, 1.3072, 1.7691, -4 / 3), |
| ): |
| t = torch.tensor(t_value, dtype=torch.float64) |
| assert float(sm.bridge_mean(x0, x1, t)) == pytest.approx(mean) |
| assert float(sm.bridge_std(t)) == pytest.approx(0.4330, abs=5e-5) |
| x = sm.bridge_sample(x0, x1, t, z) |
| assert float(x) == pytest.approx(sample, abs=5e-5) |
| assert float(sm.conditional_drift(x1, x, t)) == pytest.approx(drift, abs=5e-5) |
| assert float(sm.conditional_score(x0, x1, x, t)) == pytest.approx(-0.9238, abs=5e-5) |
| assert float(sm.flow_velocity(x0, x1, x, t)) == pytest.approx(flow, abs=5e-5) |
| assert float(sm.flow_velocity_collected(x0, x1, x, t)) == pytest.approx(flow, abs=5e-5) |
| assert float(sm.flow_coefficient(t)) == pytest.approx(coefficient) |
| assert float(sm.scaled_score(x0, x1, x, t)) == pytest.approx(-0.4, abs=5e-7) |
|
|
|
|
| def test_the_batch_plan_of_sf2m(): |
| sm = load("sf2m") |
| x, y = tensor(0.0, 2.0), tensor(1.0, 3.0) |
| uniform = tensor(0.5, 0.5) |
| plan = sm.entropic_plan(x, y, uniform, uniform, 1.0) |
| assert plan.sum(dim=0).tolist() == pytest.approx([0.5, 0.5]) |
| assert float(plan[0, 0] + plan[1, 1]) > float(plan[0, 1] + plan[1, 0]) |
|
|
|
|
| |
|
|
|
|
| def test_rate_error_is_not_symmetric(): |
| db = load("discrete_bridge") |
| assert db.scaling_cost(0.5) == pytest.approx(0.1534, abs=5e-5) |
| assert db.scaling_cost(2.0) == pytest.approx(0.3863, abs=5e-5) |
| assert float(db.rate_error(tensor(1.0)[0], tensor(1.0)[0])) == pytest.approx(0.0) |
|
|
|
|
| def test_two_state_continuous_time_bridge(): |
| db = load("discrete_bridge") |
| assert torch.matrix_exp(db.Q).flatten().tolist() == pytest.approx( |
| [0.6328, 0.3672, 0.5507, 0.4493], abs=5e-5) |
| R01 = db.endpoint_law(db.Q, db.PI_0) |
| assert R01.flatten().tolist() == pytest.approx([0.3164, 0.1836, 0.2754, 0.2246], abs=5e-5) |
| assert R01.sum(dim=0).tolist() == pytest.approx([0.5918, 0.4082], abs=5e-5) |
| f, g, plan = db.sinkhorn(R01, db.PI_0, db.PI_1) |
| assert f.tolist() == pytest.approx([1.0716, 0.9373], abs=5e-5) |
| assert g.tolist() == pytest.approx([0.3349, 1.9643], abs=5e-5) |
| assert float(g[1] / g[0]) == pytest.approx(5.865, abs=5e-4) |
| assert plan.flatten().tolist() == pytest.approx([0.1136, 0.3864, 0.0864, 0.4136], abs=5e-5) |
| assert torch.matrix_exp(db.Q / 2).flatten().tolist() == pytest.approx( |
| [0.7146, 0.2854, 0.4281, 0.5719], abs=5e-5) |
| phi = db.backward_potential(db.Q, g, 0.5) |
| assert phi.tolist() == pytest.approx([0.7999, 1.2667], abs=5e-5) |
| tilted = db.tilted_generator(db.Q, phi) |
| assert tilted.flatten().tolist() == pytest.approx( |
| [-1.5836, 1.5836, 0.9472, -0.9472], abs=5e-5) |
| assert db.backward_potential(db.Q, g, 0.0).tolist() == pytest.approx([0.9331, 1.0669], abs=5e-5) |
| at_zero = db.tilted_generator(db.Q, db.backward_potential(db.Q, g, 0.0)) |
| assert float(at_zero[0, 1]) == pytest.approx(1.1433, abs=5e-5) |
| assert float(at_zero[1, 0]) == pytest.approx(1.3120, abs=5e-5) |
| at_one = db.tilted_generator(db.Q, g) |
| assert float(at_one[0, 1]) == pytest.approx(5.8652, abs=5e-5) |
| assert float(at_one[1, 0]) == pytest.approx(0.2557, abs=5e-5) |
| marginal = db.bridge_marginal(db.Q, db.PI_0, f, g, 0.5) |
| assert marginal.tolist() == pytest.approx([0.4668, 0.5332], abs=5e-5) |
| assert db.bridge_marginal(db.Q, db.PI_0, f, g, 1.0).tolist() == pytest.approx([0.2, 0.8]) |
| reverse = db.reverse_rates(tilted, marginal) |
| assert float(reverse[1, 0]) == pytest.approx(1.3862, abs=5e-5) |
| assert float(reverse[0, 1]) == pytest.approx(1.0821, abs=5e-5) |
| assert float(marginal[0] * tilted[0, 1]) == pytest.approx( |
| float(marginal[1] * reverse[1, 0]), abs=5e-5) |
|
|
|
|
| def test_euler_sampler_and_the_cost_computed_twice(): |
| db = load("discrete_bridge") |
| R01 = db.endpoint_law(db.Q, db.PI_0) |
| f, g, plan = db.sinkhorn(R01, db.PI_0, db.PI_1) |
| errors = {} |
| for steps, expected in ((10, 0.0223), (20, 0.0106), (40, 0.0051), (100, 0.0021)): |
| terminal, leaving = db.euler_marginal(db.Q, db.PI_0, f, g, steps) |
| errors[steps] = float((terminal - db.PI_1).abs().sum() / 2) |
| assert errors[steps] == pytest.approx(expected, abs=1e-4) |
| if steps == 10: |
| assert leaving[:3] == pytest.approx([0.1143, 0.1187, 0.1245], abs=5e-5) |
| assert leaving[-1] == pytest.approx(0.3649, abs=5e-5) |
| assert db.relative_entropy(plan, R01) == pytest.approx(0.3235, abs=5e-5) |
| integral, at_half = db.path_cost(db.Q, db.PI_0, f, g) |
| assert at_half == pytest.approx(0.1300, abs=5e-5) |
| assert integral == pytest.approx(0.3235, abs=5e-5) |
| expected = {0.5: 0.3148, 1.0: 0.3235, 2.0: 0.3337, 4.0: 0.3348} |
| for lam, value in expected.items(): |
| law = db.endpoint_law(lam * db.Q, db.PI_0) |
| _, _, scaled = db.sinkhorn(law, db.PI_0, db.PI_1) |
| assert db.relative_entropy(scaled, law) == pytest.approx(value, abs=5e-5) |
| assert db.relative_entropy(db.PI_1, db.STATIONARY) == pytest.approx(0.3348, abs=5e-5) |
|
|
|
|
| |
|
|
|
|
| def test_ddsbm_single_bond(): |
| dd = load("ddsbm") |
| kernel = dd.replacement_kernel(0.4, dd.PRIOR) |
| assert kernel.flatten().tolist() == pytest.approx( |
| [0.6, 0.2, 0.2, 0.2, 0.6, 0.2, 0.2, 0.2, 0.6]) |
| reference = torch.diag(dd.BOND_0) @ kernel |
| assert reference.sum(dim=0).tolist() == pytest.approx([0.48, 0.28, 0.24]) |
| assert (dd.BOND_1 / reference.sum(dim=0)).tolist() == pytest.approx( |
| [0.2083, 0.7143, 2.9167], abs=5e-5) |
| f, g, plan, errors = dd.sinkhorn(reference, dd.BOND_0, dd.BOND_1, 200) |
| assert f.tolist() == pytest.approx([1.1779, 0.9506, 0.5062], abs=5e-5) |
| assert g.tolist() == pytest.approx([0.1842, 0.6918, 3.0004], abs=5e-5) |
| assert errors[-1] < 1e-12 |
| assert float(plan[0, 2]) == pytest.approx(0.495, abs=5e-4) |
| assert dd.held_mass(plan) == pytest.approx(0.261, abs=5e-4) |
| assert float((dd.BOND_0 * dd.BOND_1).sum()) == pytest.approx(0.18) |
| assert float(torch.minimum(dd.BOND_0, dd.BOND_1).sum()) == pytest.approx(0.4) |
| expected = {0.2: 0.2250, 0.4: 0.2612, 0.6: 0.2926, 0.8: 0.3248, 0.99: 0.3812} |
| for r, value in expected.items(): |
| held = dd.sinkhorn(torch.diag(dd.BOND_0) @ dd.replacement_kernel(r, dd.PRIOR), |
| dd.BOND_0, dd.BOND_1, 200)[2] |
| assert dd.held_mass(held) == pytest.approx(value, abs=5e-5) |
|
|
|
|
| def test_ddsbm_pinned_rates_and_the_posterior_average(): |
| dd = load("ddsbm") |
| beta = dd.resampling_rate(0.4) |
| assert beta == pytest.approx(0.9163, abs=5e-5) |
| rates = dd.replacement_generator(beta, dd.PRIOR) |
| assert float(rates[0, 1]) == pytest.approx(0.3054, abs=5e-5) |
| assert rates.sum(dim=1).abs().max() < 1e-12 |
| assert dd.half_survival_time(0.4) == pytest.approx(math.log(0.8) / math.log(0.4)) |
| assert dd.half_survival_time(0.4) == pytest.approx(0.2435, abs=5e-5) |
| pinned = dd.replacement_kernel(0.5, dd.PRIOR) |
| assert float(pinned[0, 0]) == pytest.approx(2 / 3, abs=5e-5) |
| assert float(pinned[0, 1]) == pytest.approx(1 / 6, abs=5e-5) |
| toward = dd.pinned_rate(rates, pinned, 0, 2, 2) |
| away = dd.pinned_rate(rates, pinned, 2, 0, 2) |
| between = dd.pinned_rate(rates, pinned, 0, 1, 2) |
| assert toward == pytest.approx(1.2217, abs=5e-5) |
| assert away == pytest.approx(0.0764, abs=5e-5) |
| assert between == pytest.approx(0.3054, abs=5e-5) |
| assert toward / away == pytest.approx(16.0, abs=5e-4) |
| assert dd.rate_error(toward, 1.0) == pytest.approx(0.0229, abs=5e-5) |
| posterior = 0.6 * toward + 0.4 * between |
| assert posterior == pytest.approx(0.8552, abs=5e-5) |
| rates_two = [toward, between] |
| weights = [0.6, 0.4] |
| assert dd.posterior_loss(rates_two, weights, posterior) == pytest.approx(0.1357, abs=5e-5) |
| assert dd.posterior_loss(rates_two, weights, 1.0) == pytest.approx(0.1467, abs=5e-5) |
| assert dd.posterior_loss(rates_two, weights, 0.5) == pytest.approx(0.2395, abs=5e-5) |
|
|
|
|
| def test_ddsbm_factorization_counts(): |
| dd = load("ddsbm") |
| assert dd.feature_count(9) == (9, 36, 45) |
| assert dd.state_space_size(9, 5, 4) == pytest.approx(9.2e27, rel=5e-3) |
| assert dd.total_jump_rate(36, dd.resampling_rate(0.4), 4) == pytest.approx(24.74, abs=5e-3) |
|
|
|
|
| |
|
|
|
|
| def test_csbm_references(): |
| cs = load("csbm") |
| A = cs.uniform_reference(0.35, 2) |
| assert A.flatten().tolist() == pytest.approx([0.65, 0.35, 0.35, 0.65]) |
| assert torch.linalg.matrix_power(A, 3).flatten().tolist() == pytest.approx( |
| [0.5135, 0.4865, 0.4865, 0.5135], abs=5e-5) |
| assert cs.uniform_reference(0.35, 5).sum(dim=1).tolist() == pytest.approx([1.0] * 5) |
| ordered = cs.ordered_reference(1.0, 5) |
| assert ordered[2].tolist() == pytest.approx( |
| [0.0545, 0.2442, 0.4026, 0.2442, 0.0545], abs=5e-5) |
| assert cs.transition_loss(tensor(0.3160, 0.6840), tensor(0.25, 0.75)) == pytest.approx( |
| 0.0110, abs=5e-5) |
|
|
|
|
| def test_csbm_dimf_on_sixteen_paths(): |
| cs = load("csbm") |
| A = cs.uniform_reference(cs.ALPHA, 2) |
| plan = cs.sinkhorn(torch.diag(cs.PI_0) @ torch.linalg.matrix_power(A, 3), cs.PI_0, cs.PI_1) |
| assert plan.flatten().tolist() == pytest.approx( |
| [0.10432, 0.39568, 0.09568, 0.40432], abs=5e-6) |
| assert float(cs.bridge_table(A, 3)[0, 0, 0, 0]) == pytest.approx(0.5348, abs=5e-5) |
| gap, start, target = cs.starting_gap(A, 3, cs.PI_0, cs.PI_1) |
| assert float(target[0, 0, 0, 0]) == pytest.approx(0.0558, abs=5e-5) |
| assert float(target[1, 1, 1, 1]) == pytest.approx(0.2162, abs=5e-5) |
| assert float(start[0, 0, 0, 0]) == pytest.approx(0.0535, abs=5e-5) |
| assert float(start[1, 1, 1, 1]) == pytest.approx(0.2139, abs=5e-5) |
| assert gap == pytest.approx(0.0173, abs=5e-5) |
| for axis, value in ((1, [0.4754, 0.5246]), (2, [0.4107, 0.5893])): |
| marginal = start.sum(dim=[d for d in range(4) if d != axis]) |
| assert marginal.tolist() == pytest.approx(value, abs=5e-5) |
| gaps, conditionals = cs.d_imf(cs.PI_0[:, None] * cs.PI_1[None, :], A, 3, 5, target) |
| assert conditionals[0].flatten().tolist() == pytest.approx( |
| [0.6243, 0.3757, 0.3265, 0.6735], abs=5e-5) |
| assert conditionals[1].flatten().tolist() == pytest.approx( |
| [0.5626, 0.4374, 0.2729, 0.7271], abs=5e-5) |
| assert conditionals[2].flatten().tolist() == pytest.approx( |
| [0.3160, 0.6840, 0.1192, 0.8808], abs=5e-5) |
| assert gaps == pytest.approx([2.9e-4, 5.0e-6, 8.4e-8, 1.4e-9, 2.4e-11], rel=0.05) |
| assert gaps[0] / gaps[1] == pytest.approx(59.0, abs=1.0) |
|
|
|
|
| def test_csbm_mixing_and_length(): |
| cs = load("csbm") |
| independent = cs.PI_0[:, None] * cs.PI_1[None, :] |
| expected = {0.2: (0.13589, 5.0), 0.4: (0.00512, 284.3)} |
| for alpha, (start, ratio) in expected.items(): |
| A = cs.uniform_reference(alpha, 2) |
| gap, _, target = cs.starting_gap(A, 3, cs.PI_0, cs.PI_1) |
| assert gap == pytest.approx(start, abs=5e-5) |
| later, _ = cs.d_imf(independent, A, 3, 2, target) |
| assert later[0] / later[1] == pytest.approx(ratio, rel=0.01) |
| A = cs.uniform_reference(cs.ALPHA, 2) |
| for steps, (start, ratio) in {2: (0.0574, 7.3), 5: (0.00156, 4674.0)}.items(): |
| gap, _, target = cs.starting_gap(A, steps, cs.PI_0, cs.PI_1) |
| assert gap == pytest.approx(start, abs=5e-5) |
| later, _ = cs.d_imf(independent, A, steps, 2, target) |
| assert later[0] / later[1] == pytest.approx(ratio, rel=0.01) |
|
|
|
|
| |
|
|
|
|
| def test_two_branches_worked_out(): |
| br = load("branched") |
| linear = lambda t: 0.3 * t |
| late = lambda t: 0.3 * t ** 2 |
| assert float(br.branch_weights(linear, torch.tensor(0.5))) == pytest.approx(0.15) |
| assert float(br.branch_weights(linear, torch.tensor(1.0))) == pytest.approx(0.3) |
| near_zero = br.weighted_energy(2.0, lambda t: 1.0 - linear(t)) |
| near_one = br.weighted_energy(-2.0, linear) |
| assert near_zero == pytest.approx(1.7, abs=5e-5) |
| assert near_one == pytest.approx(0.3, abs=5e-5) |
| assert near_zero + near_one == pytest.approx(2.0, abs=5e-5) |
| far_linear = near_zero + br.weighted_energy(-4.0, linear) |
| far_late = br.weighted_energy(2.0, lambda t: 1.0 - late(t)) + br.weighted_energy(-4.0, late) |
| assert far_linear == pytest.approx(2.9, abs=5e-5) |
| assert far_late == pytest.approx(2.6, abs=5e-5) |
|
|
|
|
| def test_the_interpolant(): |
| br = load("branched") |
| grid = torch.tensor([0.0, 0.5, 1.0]) |
| assert br.interpolant(0.0, 2.0, 1.0, grid).tolist() == pytest.approx([0.0, 1.25, 2.0]) |
| assert br.interpolant_velocity(0.0, 2.0, 1.0, grid).tolist() == pytest.approx([3.0, 2.0, 1.0]) |
| assert br.interpolant_energy(0.0, 2.0, 1.0) == pytest.approx(2.1667, abs=5e-5) |
| assert br.interpolant_energy(0.0, 2.0, 0.0) == pytest.approx(2.0, abs=5e-5) |
|
|
|
|
| def test_entangled_tilt_and_parameterization(): |
| br = load("branched") |
| normalizer = 0.1 * 1.0 + 0.9 * 0.01 |
| assert normalizer == pytest.approx(0.109) |
| tilted = 0.1 / normalizer |
| assert tilted == pytest.approx(0.917, abs=5e-4) |
| assert tilted / 0.1 == pytest.approx(9.17, abs=5e-3) |
| reward = (1.0 - tilted) * math.log(0.01) |
| assert reward == pytest.approx(-0.380, abs=5e-4) |
| assert reward - math.log(normalizer) == pytest.approx(1.836, abs=5e-4) |
| assert -math.log(normalizer) == pytest.approx(2.216, abs=5e-4) |
| assert br.NOISE == pytest.approx(1.4142, abs=5e-5) |
| step = (-0.3 + 0.5) * 0.01 - 1.0 * 0.8 * 0.01 + br.NOISE * math.sqrt(0.01) * 0.4 |
| assert step == pytest.approx(0.0506, abs=5e-5) |
| assert 0.5 / br.NOISE == pytest.approx(0.3536, abs=5e-5) |
| assert 0.5 * (0.5 / br.NOISE) ** 2 * 0.01 == pytest.approx(0.000625, abs=5e-7) |
| force = br.compose_bias(torch.tensor([[0.5]]), torch.tensor([[[0.3, -0.4]]]), |
| torch.tensor([[[1.0, 0.0]]]))[0, 0] |
| assert force.tolist() == pytest.approx([0.5, -0.4]) |
| assert float(force.norm()) == pytest.approx(0.640, abs=5e-4) |
| weights = torch.softmax(torch.tensor([-1.0, 0.0, 1.0]), dim=0) |
| assert weights.tolist() == pytest.approx([0.090, 0.245, 0.665], abs=5e-4) |
|
|
|
|
| def test_entangled_training_lowers_the_loss(): |
| br = load("branched") |
| torch.manual_seed(0) |
| generator = torch.Generator().manual_seed(0) |
| force = br.BiasForce(width=16) |
| rollout = br.roll_out(force, 16, 8, torch.device("cpu"), generator) |
| weights = br.importance_weights(rollout) |
| assert float(weights.sum()) == pytest.approx(1.0) |
| assert br.force_check(rollout["positions"][-1]) < 1e-5 |
| loss = br.cross_entropy_loss(force, rollout, weights) |
| assert torch.isfinite(loss) |
|
|
|
|
| |
|
|
|
|
| @pytest.mark.parametrize("name,flags", [ |
| ("optimal_transport", ["--steps", "300"]), |
| ("sinkhorn", ["--steps", "20"]), |
| ("displacement", ["--steps", "64"]), |
| ("ot_flow_matching", ["--steps", "3", "--samples", "64", "--sample-steps", "5"]), |
| ("static_bridge", ["--steps", "3"]), |
| ("dynamic_bridge", ["--steps", "20"]), |
| ("dsb", ["--steps", "2"]), |
| ("dsbm", ["--steps", "3"]), |
| ("sf2m", ["--steps", "21"]), |
| ("discrete_bridge", ["--steps", "4"]), |
| ("ddsbm", ["--steps", "20"]), |
| ("csbm", ["--steps", "3"]), |
| ("branched", ["--steps", "2", "--samples", "16"]), |
| ]) |
| def test_the_script_runs(name, flags, monkeypatch, capsys): |
| module = load(name) |
| monkeypatch.setattr(sys, "argv", [name] + flags + ["--quiet"]) |
| module.main() |
| assert capsys.readouterr().out.strip() |
|
|
