Skip to content

Commit 1c35fc7

Browse files
jessegrabowskijessegrabowski
committed
Add tests for names/dims/coords
1 parent 3204a19 commit 1c35fc7

File tree

2 files changed

+315
-0
lines changed

2 files changed

+315
-0
lines changed

‎pymc_extras/statespace/models/DFM.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -238,6 +238,8 @@ def __init__(
238238
verbose: bool = True,
239239
):
240240
"""
241+
Create a Bayesian Dynamic Factor Model.
242+
241243
Parameters
242244
----------
243245
k_factors : int

‎tests/statespace/models/test_DFM.py

Lines changed: 313 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,16 @@
1414

1515
from pymc_extras.statespace.models.DFM import BayesianDynamicFactor
1616
from pymc_extras.statespace.utils.constants import (
17+
ALL_STATE_AUX_DIM,
18+
ALL_STATE_DIM,
19+
AR_PARAM_DIM,
20+
ERROR_AR_PARAM_DIM,
21+
EXOG_STATE_DIM,
22+
FACTOR_DIM,
1723
LONG_MATRIX_NAMES,
1824
MATRIX_NAMES,
25+
OBS_STATE_AUX_DIM,
26+
OBS_STATE_DIM,
1927
SHORT_NAME_TO_LONG,
2028
)
2129
from tests.statespace.shared_fixtures import rng
@@ -367,3 +375,308 @@ def test_DFM_exog_shared_vs_not(shared):
367375
assert not np.allclose(
368376
contributions[0], contributions[1]
369377
), f"Expected different contributions, got {contributions}"
378+
379+
380+
class TestDFMConfiguration:
381+
def test_static_factor_no_ar_no_exog_diagonal_error(self):
382+
mod = BayesianDynamicFactor(
383+
k_factors=1,
384+
factor_order=0,
385+
k_endog=3,
386+
endog_names=["y0", "y1", "y2"],
387+
error_order=0,
388+
error_var=False,
389+
error_cov_type="diagonal",
390+
measurement_error=False,
391+
verbose=False,
392+
)
393+
394+
expected_param_names = ["x0", "P0", "factor_loadings", "error_sigma"]
395+
expected_param_dims = {
396+
"x0": (ALL_STATE_DIM,),
397+
"P0": (ALL_STATE_DIM, ALL_STATE_AUX_DIM),
398+
"factor_loadings": (OBS_STATE_DIM, FACTOR_DIM),
399+
"error_sigma": (OBS_STATE_DIM,),
400+
}
401+
expected_coords = {
402+
OBS_STATE_DIM: ["y0", "y1", "y2"],
403+
ALL_STATE_DIM: ["L0.factor_0"],
404+
ALL_STATE_AUX_DIM: ["L0.factor_0"],
405+
FACTOR_DIM: ["factor_1"],
406+
}
407+
408+
assert mod.param_names == expected_param_names
409+
assert mod.param_dims == expected_param_dims
410+
for k, v in expected_coords.items():
411+
assert mod.coords[k] == v
412+
assert mod.state_names == ["L0.factor_0"]
413+
assert mod.observed_states == ["y0", "y1", "y2"]
414+
assert mod.shock_names == ["factor_shock_0"]
415+
416+
def test_dynamic_factor_ar1_error_diagonal_error(self):
417+
mod = BayesianDynamicFactor(
418+
k_factors=2,
419+
factor_order=2,
420+
k_endog=3,
421+
endog_names=["y0", "y1", "y2"],
422+
error_order=1,
423+
error_var=False,
424+
error_cov_type="diagonal",
425+
measurement_error=True,
426+
verbose=False,
427+
)
428+
expected_param_names = [
429+
"x0",
430+
"P0",
431+
"factor_loadings",
432+
"factor_ar",
433+
"error_ar",
434+
"error_sigma",
435+
"sigma_obs",
436+
]
437+
expected_param_dims = {
438+
"x0": (ALL_STATE_DIM,),
439+
"P0": (ALL_STATE_DIM, ALL_STATE_AUX_DIM),
440+
"factor_loadings": (OBS_STATE_DIM, FACTOR_DIM),
441+
"factor_ar": (FACTOR_DIM, AR_PARAM_DIM),
442+
"error_ar": (OBS_STATE_DIM, ERROR_AR_PARAM_DIM),
443+
"error_sigma": (OBS_STATE_DIM,),
444+
"sigma_obs": (OBS_STATE_DIM,),
445+
}
446+
expected_coords = {
447+
OBS_STATE_DIM: ["y0", "y1", "y2"],
448+
ALL_STATE_DIM: [
449+
"L0.factor_0",
450+
"L1.factor_0",
451+
"L0.factor_1",
452+
"L1.factor_1",
453+
"L0.error_0",
454+
"L0.error_1",
455+
"L0.error_2",
456+
],
457+
ALL_STATE_AUX_DIM: [
458+
"L0.factor_0",
459+
"L1.factor_0",
460+
"L0.factor_1",
461+
"L1.factor_1",
462+
"L0.error_0",
463+
"L0.error_1",
464+
"L0.error_2",
465+
],
466+
FACTOR_DIM: ["factor_1", "factor_2"],
467+
AR_PARAM_DIM: list(range(1, 2 * 2 + 1)),
468+
ERROR_AR_PARAM_DIM: [1],
469+
}
470+
471+
assert mod.param_names == expected_param_names
472+
assert mod.param_dims == expected_param_dims
473+
for k, v in expected_coords.items():
474+
assert mod.coords[k] == v
475+
assert len(mod.state_names) == 7
476+
assert mod.observed_states == ["y0", "y1", "y2"]
477+
assert len(mod.shock_names) == 5
478+
479+
def test_dynamic_factor_ar2_error_var_unstructured(self):
480+
mod = BayesianDynamicFactor(
481+
k_factors=1,
482+
factor_order=1,
483+
k_endog=3,
484+
endog_names=["y0", "y1", "y2"],
485+
error_order=2,
486+
error_var=True,
487+
error_cov_type="unstructured",
488+
measurement_error=True,
489+
verbose=False,
490+
)
491+
expected_param_names = [
492+
"x0",
493+
"P0",
494+
"factor_loadings",
495+
"factor_ar",
496+
"error_ar",
497+
"error_cov",
498+
"sigma_obs",
499+
]
500+
expected_param_dims = {
501+
"x0": (ALL_STATE_DIM,),
502+
"P0": (ALL_STATE_DIM, ALL_STATE_AUX_DIM),
503+
"factor_loadings": (OBS_STATE_DIM, FACTOR_DIM),
504+
"factor_ar": (FACTOR_DIM, AR_PARAM_DIM),
505+
"error_ar": (OBS_STATE_DIM, ERROR_AR_PARAM_DIM),
506+
"error_cov": (OBS_STATE_DIM, OBS_STATE_AUX_DIM),
507+
"sigma_obs": (OBS_STATE_DIM,),
508+
}
509+
expected_coords = {
510+
OBS_STATE_DIM: ["y0", "y1", "y2"],
511+
ALL_STATE_DIM: [
512+
"L0.factor_0",
513+
"L0.error_0",
514+
"L1.error_0",
515+
"L0.error_1",
516+
"L1.error_1",
517+
"L0.error_2",
518+
"L1.error_2",
519+
],
520+
ALL_STATE_AUX_DIM: [
521+
"L0.factor_0",
522+
"L0.error_0",
523+
"L1.error_0",
524+
"L0.error_1",
525+
"L1.error_1",
526+
"L0.error_2",
527+
"L1.error_2",
528+
],
529+
FACTOR_DIM: ["factor_1"],
530+
AR_PARAM_DIM: [1],
531+
ERROR_AR_PARAM_DIM: list(range(1, 2 * 3 + 1)),
532+
}
533+
534+
assert mod.param_names == expected_param_names
535+
assert mod.param_dims == expected_param_dims
536+
for k, v in expected_coords.items():
537+
assert mod.coords[k] == v
538+
assert len(mod.state_names) == 7
539+
assert mod.observed_states == ["y0", "y1", "y2"]
540+
assert len(mod.shock_names) == 7
541+
542+
def test_exog_shared_exog_states_exog_innovations(self):
543+
mod = BayesianDynamicFactor(
544+
k_factors=2,
545+
factor_order=1,
546+
k_endog=3,
547+
endog_names=["y0", "y1", "y2"],
548+
error_order=1,
549+
error_var=False,
550+
k_exog=2,
551+
exog_names=["x0", "x1"],
552+
shared_exog_states=True,
553+
exog_innovations=True,
554+
error_cov_type="diagonal",
555+
measurement_error=True,
556+
verbose=False,
557+
)
558+
expected_param_names = [
559+
"x0",
560+
"P0",
561+
"factor_loadings",
562+
"factor_ar",
563+
"error_ar",
564+
"error_sigma",
565+
"sigma_obs",
566+
"beta",
567+
"beta_sigma",
568+
]
569+
expected_param_dims = {
570+
"x0": (ALL_STATE_DIM,),
571+
"P0": (ALL_STATE_DIM, ALL_STATE_AUX_DIM),
572+
"factor_loadings": (OBS_STATE_DIM, FACTOR_DIM),
573+
"factor_ar": (FACTOR_DIM, AR_PARAM_DIM),
574+
"error_ar": (OBS_STATE_DIM, ERROR_AR_PARAM_DIM),
575+
"error_sigma": (OBS_STATE_DIM,),
576+
"sigma_obs": (OBS_STATE_DIM,),
577+
"beta": (EXOG_STATE_DIM,),
578+
"beta_sigma": (EXOG_STATE_DIM,),
579+
}
580+
expected_coords = {
581+
OBS_STATE_DIM: ["y0", "y1", "y2"],
582+
ALL_STATE_DIM: [
583+
"L0.factor_0",
584+
"L0.factor_1",
585+
"L0.error_0",
586+
"L0.error_1",
587+
"L0.error_2",
588+
"beta_x0[shared]",
589+
"beta_x1[shared]",
590+
],
591+
ALL_STATE_AUX_DIM: [
592+
"L0.factor_0",
593+
"L0.factor_1",
594+
"L0.error_0",
595+
"L0.error_1",
596+
"L0.error_2",
597+
"beta_x0[shared]",
598+
"beta_x1[shared]",
599+
],
600+
FACTOR_DIM: ["factor_1", "factor_2"],
601+
AR_PARAM_DIM: [1, 2],
602+
ERROR_AR_PARAM_DIM: [1],
603+
EXOG_STATE_DIM: [1, 2],
604+
}
605+
606+
assert mod.param_names == expected_param_names
607+
assert mod.param_dims == expected_param_dims
608+
for k, v in expected_coords.items():
609+
assert mod.coords[k] == v
610+
assert len(mod.state_names) == 7
611+
assert mod.observed_states == ["y0", "y1", "y2"]
612+
assert len(mod.shock_names) == 7
613+
614+
def test_exog_not_shared_no_exog_innovations(self):
615+
mod = BayesianDynamicFactor(
616+
k_factors=1,
617+
factor_order=2,
618+
k_endog=3,
619+
endog_names=["y0", "y1", "y2"],
620+
error_order=1,
621+
error_var=False,
622+
k_exog=1,
623+
exog_names=["x0"],
624+
shared_exog_states=False,
625+
exog_innovations=False,
626+
error_cov_type="scalar",
627+
measurement_error=False,
628+
verbose=False,
629+
)
630+
expected_param_names = [
631+
"x0",
632+
"P0",
633+
"factor_loadings",
634+
"factor_ar",
635+
"error_ar",
636+
"error_sigma",
637+
"beta",
638+
]
639+
expected_param_dims = {
640+
"x0": (ALL_STATE_DIM,),
641+
"P0": (ALL_STATE_DIM, ALL_STATE_AUX_DIM),
642+
"factor_loadings": (OBS_STATE_DIM, FACTOR_DIM),
643+
"factor_ar": (FACTOR_DIM, AR_PARAM_DIM),
644+
"error_ar": (OBS_STATE_DIM, ERROR_AR_PARAM_DIM),
645+
"error_sigma": (),
646+
"beta": (EXOG_STATE_DIM,),
647+
}
648+
expected_coords = {
649+
OBS_STATE_DIM: ["y0", "y1", "y2"],
650+
ALL_STATE_DIM: [
651+
"L0.factor_0",
652+
"L1.factor_0",
653+
"L0.error_0",
654+
"L0.error_1",
655+
"L0.error_2",
656+
"beta_x0[y0]",
657+
"beta_x0[y1]",
658+
"beta_x0[y2]",
659+
],
660+
ALL_STATE_AUX_DIM: [
661+
"L0.factor_0",
662+
"L1.factor_0",
663+
"L0.error_0",
664+
"L0.error_1",
665+
"L0.error_2",
666+
"beta_x0[y0]",
667+
"beta_x0[y1]",
668+
"beta_x0[y2]",
669+
],
670+
FACTOR_DIM: ["factor_1"],
671+
AR_PARAM_DIM: [1, 2],
672+
ERROR_AR_PARAM_DIM: [1],
673+
EXOG_STATE_DIM: [1, 2, 3],
674+
}
675+
676+
assert mod.param_names == expected_param_names
677+
assert mod.param_dims == expected_param_dims
678+
for k, v in expected_coords.items():
679+
assert mod.coords[k] == v
680+
assert len(mod.state_names) == 8
681+
assert mod.observed_states == ["y0", "y1", "y2"]
682+
assert len(mod.shock_names) == 7

0 commit comments

Comments
 (0)