1import numpy as np2import torch34from chronos_x.models.joint_path import samples_from_copula, LowRankCopulaHead, pit_student_t, z_from_pit567def test_copula_samples_reproduce_marginals():8 torch.manual_seed(0)9 n, B, H = 20000, 2, 6010 z = torch.randn(n, B, H)11 mu = torch.full((B, H), 0.001); sc = torch.full((B, H), 0.002); nu = torch.full((B, H), 6.0)12 S = samples_from_copula(z, mu, sc, nu)13 assert abs(float(S.mean()) - 0.001) < 5e-514 assert abs(float(S.std()) - 0.002 * np.sqrt(6 / 4)) < 1e-4151617def test_sample_z_is_standardised_and_correlation_consistent():18 torch.manual_seed(1)19 head = LowRankCopulaHead(8, 60, rank=2)20 with torch.no_grad():21 head.load.weight.mul_(20.0) # strong common factor22 h = torch.randn(3, 60, 8)23 z = head.sample_z(h, 20000)24 R = head.correlation(h)25 emp = torch.corrcoef(z[:, 0, :].T)26 assert abs(float(z.std()) - 1) < 0.0227 assert float((emp - R[0]).abs().mean()) < 0.03282930def test_copula_nll_minimised_at_independence_for_independent_data():31 torch.manual_seed(2)32 head = LowRankCopulaHead(8, 60, rank=2)33 h = torch.randn(64, 60, 8)34 z = torch.randn(64, 60); mask = torch.ones(64, 60, dtype=torch.bool)35 nll_small = head.nll(h, z, mask)36 with torch.no_grad():37 head.load.weight.mul_(30.0)38 nll_big = head.nll(h, z, mask)39 assert nll_small < nll_big404142def test_pit_roundtrip():43 y = np.array([[-0.003, 0.0, 0.004]]); mu = np.zeros((1, 3)); sc = np.full((1, 3), 0.002); nu = np.full((1, 3), 5.0)44 u = pit_student_t(y, mu, sc, nu)45 assert 0 < u[0, 0] < 0.5 < u[0, 2] < 1 and abs(u[0, 1] - 0.5) < 1e-1246 z = z_from_pit(u)47 assert z[0, 0] < 0 < z[0, 2]48