Merge pull request #76 from francescoTheSantis/mid_cleaning

Mid cleaning

Author: pietrobarbiero

doc/guides/using_mid_level.rst

@@ -183,6 +183,12 @@ Concept Bottleneck Model ``input → latent → concepts → task`` as a probabi
       point predictions.
     - :class:`~torch_concepts.nn.AncestralSamplingInference` — draws a (reparameterised) sample per
       variable in topological order.
+    - :class:`~torch_concepts.nn.BeliefPropagation` — differentiable sum-product marginals over any
+      factor graph (directed, undirected or mixed). Exact on trees, the standard approximation with
+      loops. This is the engine to *train* an undirected model with.
+    - :class:`~torch_concepts.nn.PgmpyVariableElimination` — **exact** marginals on any factor
+      graph, via pgmpy. Evaluation only: it exports a static table per observation, so it carries
+      no gradients.
     - :class:`~torch_concepts.nn.ForwardInference`, :class:`~torch_concepts.nn.IndependentInference`,
       :class:`~torch_concepts.nn.RejectionSampling`, :class:`~torch_concepts.nn.ImportanceSampling`,
       and the Pyro-backed :class:`~torch_concepts.nn.VariationalInference` provide further

doc/modules/mid_level_api.rst

@@ -58,7 +58,6 @@ Probabilistic Models
    ProbabilisticModel
    BayesianNetwork
    MarkovNetwork
-   ChainGraph
 
 Inference
 ---------
@@ -73,10 +72,10 @@ Inference
    AncestralSamplingInference
    MAPForwardInference
    BeliefPropagation
+   PgmpyVariableElimination
    RejectionSampling
    ImportanceSampling
    VariationalInference
-   PyroImportanceSampling
    BaseProposal
    MutilatedNetworkProposal
 
@@ -90,3 +89,4 @@ Base Classes
    BaseInference
    TorchBaseInference
    PyroBaseInference
+   PgmpyBaseInference

examples/experiments/plate_refactor_acceptance.py

@@ -177,18 +177,13 @@ def verify():
     # 9. (post-refactor) Pyro engines, skipped if pyro is absent
     try:
         import pyro  # noqa: F401
-        from torch_concepts.nn import VariationalInference, PyroImportanceSampling
+        from torch_concepts.nn import VariationalInference
         vi = VariationalInference(pgm)
         seed_everything(0)
         y_hi = vi.query(query=["y"], evidence={"x": xt, "c1": ones}).probs["y"]
         seed_everything(0)
         y_lo = vi.query(query=["y"], evidence={"x": xt, "c1": zeros}).probs["y"]
         assert not torch.allclose(y_hi, y_lo), "member evidence ignored by model_fn"
-        with warnings.catch_warnings():
-            warnings.simplefilter("ignore")
-            pis = PyroImportanceSampling(pgm, n_samples=200)
-            p = pis.query({"y": ones4}, evidence={"c1": ones4}).probabilities
-            assert p.shape == (B4,) and torch.isfinite(p).all()
         print("9. Pyro engines honor member evidence                     OK")
     except ImportError:
         print("9. pyro-ppl not installed                                 SKIPPED")

examples/utilization/1_pgm/2_test_time_inference.py

@@ -26,23 +26,23 @@
 from torch_concepts.data import BnLearnDataset
 from torch_concepts.nn import ParametricCPD, BayesianNetwork, \
     AncestralSamplingInference, LearnablePrior, RejectionSampling, \
-    ImportanceSampling, MutilatedNetworkProposal, PyroImportanceSampling, \
-    BeliefPropagation
+    ImportanceSampling, MutilatedNetworkProposal, BeliefPropagation, \
+    PgmpyVariableElimination
 
 
 def _compare_query(
     query_name, evidence_names, targets,
-    rejection_engine, torch_is_engine, pyro_is_engine, bp_engine,
+    rejection_engine, is_engine, bp_engine, ve_engine,
 ):
     """Estimate P(query_name=1 | evidence) with every engine and the data.
 
     Enumerates every {0,1} assignment of ``evidence_names`` (each a single
-    binary node), queries the four inference engines plus the empirical
+    binary node), queries every inference engine plus the empirical
     frequency in ``targets``, and prints one row per estimator.
     """
     combos = list(itertools.product([0., 1.], repeat=len(evidence_names)))
-    rows = {"empirical": [], "reject S": [], "torch IS": [], "pyro  IS": [],
-            "belief prop": []}
+    rows = {"empirical": [], "reject S": [], "torch IS": [],
+            "belief prop": [], "exact VE": []}
 
     for combo in combos:
         evidence = {n: torch.tensor([[v]]) for n, v in zip(evidence_names, combo)}
@@ -51,12 +51,13 @@ def _compare_query(
         rows["reject S"].append(
             rejection_engine.query(query, evidence).probabilities.item())
         rows["torch IS"].append(
-            torch_is_engine.query(query, evidence).probabilities.item())
-        rows["pyro  IS"].append(
-            pyro_is_engine.query(query, evidence).probabilities.item())
+            is_engine.query(query, evidence).probabilities.item())
         rows["belief prop"].append(
             bp_engine.query(query=[query_name], evidence=evidence)
             .probs[query_name].item())
+        rows["exact VE"].append(
+            ve_engine.query(query=[query_name], evidence=evidence)
+            .probs[query_name].item())
 
         mask = torch.ones(targets[query_name].shape[0], dtype=torch.bool)
         for n, v in zip(evidence_names, combo):
@@ -71,8 +72,8 @@ def _compare_query(
     fmt = lambda vals: "  ".join(f"{v:6.3f}" for v in vals)
     print(f"\n=== P({query_name}=1 | {', '.join(evidence_names)}) ===")
     print(f"{'':>14} | {header}")
-    for label in ("empirical", "reject S", "torch IS", "pyro  IS", "belief prop"):
-        print(f"{label:>14} | {fmt(rows[label])}")
+    for label, vals in rows.items():
+        print(f"{label:>14} | {fmt(vals)}")
 
 
 def main():
@@ -229,12 +230,14 @@ def main():
     # a purely forward computation.
     n_test_samples = 20_000
     rejection_engine = RejectionSampling(concept_model, n_samples=n_test_samples)
-    torch_is_engine = ImportanceSampling(
+    is_engine = ImportanceSampling(
         concept_model, MutilatedNetworkProposal(concept_model),
         n_samples=n_test_samples, initial_temperature=0.1,
     )
-    pyro_is_engine = PyroImportanceSampling(concept_model, n_samples=n_test_samples)
     bp_engine = BeliefPropagation(concept_model, iters=20, damping=0.2)
+    # Exact marginals: the reference the approximate engines are judged
+    # against. Evaluation only -- it carries no gradients.
+    ve_engine = PgmpyVariableElimination(concept_model)
 
     for query_name, evidence_names in [
         ("tub", ("either", "lung")),
@@ -243,7 +246,7 @@ def main():
     ]:
         _compare_query(
             query_name, evidence_names, targets,
-            rejection_engine, torch_is_engine, pyro_is_engine, bp_engine,
+            rejection_engine, is_engine, bp_engine, ve_engine,
         )
 
     return

tests/distributions/test_delta.py

@@ -202,9 +202,9 @@ def test_dtype_preservation(self):
         self.assertEqual(dist_float64.mean.dtype, torch.float64)
 
     def test_batch_shape(self):
-        """Test batch_shape attribute."""
+        """batch_shape follows the value, so Delta composes with Independent."""
         dist = Delta([1.0, 2.0])
-        self.assertEqual(dist.batch_shape, torch.Size([]))
+        self.assertEqual(dist.batch_shape, torch.Size([2]))
 
     def test_multiple_samples_consistency(self):
         """Test that multiple samples are consistent."""

tests/nn/modules/high/models/test_cbvae.py

@@ -90,9 +90,9 @@ def test_categorical_concepts_normalise_per_concept(self, categorical_annotation
     def test_a_categorical_plates_cpd_emits_one_block_per_member(self):
         """The plate's concept CPD emits one score block per member.
 
-        Asserted on the CPD's own output: the emitted width is the plate's
-        flattened width, laid out member-major so the distribution can normalise
-        each member's block independently.
+        Asserted on the CPD's own output, which is in member layout
+        ``(batch, n_members, states)``, so the distribution can normalise each
+        member's block independently.
         """
         # Same cardinality on both concepts, so they can share one plate.
         annotations = Annotations(
@@ -110,8 +110,8 @@ def test_a_categorical_plates_cpd_emits_one_block_per_member(self):
         model(query=list(model.pgm.variables), input=torch.rand(6, INPUT_SIZE))
 
         logits = emitted["logits"]
-        assert logits.shape == (6, 6)  # 2 members x 3 states, member-major
-        probs = logits.reshape(6, 2, 3).softmax(-1)
+        assert logits.shape == (6, 2, 3)  # 2 members x 3 states
+        probs = logits.softmax(-1)
         assert torch.allclose(probs.sum(-1), torch.ones(6, 2), atol=1e-5)
 
     def test_the_latent_prior_and_guide_produce_a_positive_scale(self, binary_annotations):

tests/nn/modules/high/models/test_cvae.py

@@ -92,8 +92,9 @@ def test_binary_concepts_report_logits(self, binary_annotations):
     def test_categorical_marginals_normalise_per_concept(self, categorical_annotations):
         model = build_model(categorical_annotations, plate=False)
         for name, cardinality in (("digit", 4), ("color", 3)):
+            # A CPD reports the member layout: one member of `cardinality` states.
             logits = model.pgm.factors[name](parent_values={})["logits"]
-            assert logits.shape == (cardinality,)
+            assert logits.shape == (1, cardinality)
             assert torch.allclose(logits.softmax(-1).sum(), torch.ones(()))
 
     def test_a_categorical_plates_marginal_normalises_each_member(self):
@@ -111,7 +112,8 @@ def test_a_categorical_plates_marginal_normalises_each_member(self):
         model = build_model(annotations, plate=True)
         assert "concepts" in model.pgm.variables  # one plate, both members
 
-        assert model.pgm.factors["concepts"](parent_values={})["logits"].shape == (8,)
+        # Member layout on the CPD (2 members x 4 states); flat on the engine output.
+        assert model.pgm.factors["concepts"](parent_values={})["logits"].shape == (2, 4)
         drawn = AncestralSamplingInference(model.pgm).query(
             query=["concepts"], evidence={}, n_samples=5
         ).samples["concepts"]
@@ -534,7 +536,8 @@ def test_a_multivariate_normal_concept_sizes_its_cholesky_factor(self):
             plate=False,
             variable_distributions={"continuous": MultivariateNormal},
         )
+        # Member layout: one member, whose `scale_tril` carries the extra rank.
         params = model.pgm.factors["v"](parent_values={})
-        assert params["loc"].shape == (3,)
-        assert params["scale_tril"].shape == (3, 3)
-        assert bool((params["scale_tril"].diagonal() > 0).all())
+        assert params["loc"].shape == (1, 3)
+        assert params["scale_tril"].shape == (1, 3, 3)
+        assert bool((params["scale_tril"].diagonal(dim1=-2, dim2=-1) > 0).all())

tests/nn/modules/mid/inference/test_assembly_caching.py

@@ -168,7 +168,9 @@ def test_samples_assembly_is_cached_consistently():
     """_assemble_samples shares the cached chunks with _assemble_params but
     keys its annotation separately, so both must stay correct."""
     eng = _engine()
-    per_variable = {"concepts": torch.rand(3, 3), "y": torch.randn(3, 2)}
+    # Realisations arrive in member layout (B, n_members, member_size), as the
+    # engines cache them.
+    per_variable = {"concepts": torch.rand(3, 3, 1), "y": torch.randn(3, 1, 2)}
     for _ in range(2):
         samples = eng._assemble_samples(per_variable, ["concepts", "y"])
         assert list(samples.annotation.labels) == ["c1", "c2", "c3", "y"]

tests/nn/modules/mid/inference/test_belief_propagation.py

@@ -11,7 +11,10 @@
 from torch_concepts.nn.modules.mid.variable import ConceptVariable
 from torch_concepts.nn.modules.mid.factors.cpd import ParametricCPD
 from torch_concepts.nn.modules.mid.factors.potential import ParametricPotential
-from torch_concepts.nn.modules.mid.inference.utils import enumerable_cardinality
+from torch_concepts.nn.modules.mid.inference.utils import (
+    enumerable_cardinality,
+    factor_table,
+)
 from torch_concepts.nn.modules.mid.graph.probabilistic_model import ProbabilisticModel
 from torch_concepts.nn.modules.mid.graph.markov_network import MarkovNetwork
 from torch_concepts.nn.modules.mid.inference.torch.belief_propagation import BeliefPropagation
@@ -357,8 +360,7 @@ def test_cells_get_independent_masks(self):
             pot.parametrization["energy"][0].weight.zero_()
         fg = MarkovNetwork(variables=[a, b], factors=[pot])
         fg.train()
-        eng = BeliefPropagation(fg, iters=3)
-        table = eng._factor_table(
+        table = factor_table(
             pot, ["a", "b"], {"a": a, "b": b}, {}, torch.Size([1]),
             torch.float32, torch.device("cpu"),
         )
@@ -545,7 +547,7 @@ def _member_probs(self, fg, out, owner_name, member):
         """The engine's marginal for one member, widened to a state distribution."""
         owner = fg.variables[owner_name]
         return _state_marginal(
-            owner.member(member), out.probs[owner_name][..., owner.column_of(member)]
+            owner.member(member), out.probs[member]
         )
 
     # -- test-time queries --------------------------------------------------
@@ -571,9 +573,9 @@ def test_partial_plate_evidence_matches_exact(self):
             assert torch.allclose(bp, exact[member], atol=1e-5), (member, bp, exact[member])
         # An observed member reports its evidence: conditioning makes it a point mass.
         g = fg.variables["g"]
-        assert torch.allclose(out.probs["g"][..., g.column_of("g1")], torch.ones(1, 1))
+        assert torch.allclose(out.probs["g1"], torch.ones(1, 1))
         h = fg.variables["h"]
-        assert torch.allclose(out.probs["h"][..., h.column_of("h3")], torch.zeros(1, 1))
+        assert torch.allclose(out.probs["h3"], torch.zeros(1, 1))
 
     def test_evidence_on_a_member_moves_the_other_plate(self):
         """Both plates hang off ``root``, so conditioning one must move the other."""
@@ -651,7 +653,7 @@ def losses():
             out = eng.query(query=["g", "h"], evidence=evidence)
             return sum(
                 _binary_ce(
-                    out.probs[owner][..., fg.variables[owner].column_of(m)], target
+                    out.probs[m], target
                 )
                 for owner, m, target in [
                     ("g", "g2", targets["g2"]),
@@ -669,9 +671,9 @@ def losses():
         assert loss.item() < first * 0.5
 
         out = eng.query(query=["g", "h"], evidence=evidence)
-        assert out.probs["g"][0, fg.variables["g"].column_of("g2")].item() > 0.9
-        assert out.probs["h"][0, fg.variables["h"].column_of("h2")].item() > 0.9
-        assert out.probs["h"][0, fg.variables["h"].column_of("h3")].item() < 0.1
+        assert out.probs["g2"][0].item() > 0.9
+        assert out.probs["h2"][0].item() > 0.9
+        assert out.probs["h3"][0].item() < 0.1
 
     # -- categorical plates (needs build_distribution's per-member split) ----
     def _mixed_plate_bn(self, seed=32):
@@ -706,13 +708,13 @@ def test_categorical_plate_members_are_independent(self):
         out = BeliefPropagation(fg, iters=30).query(query=["h"], evidence={})
         h = fg.variables["h"]
         for member in h.members:
-            block = out.probs["h"][..., h.column_of(member)]
+            block = out.probs[member]
             assert block.shape == (1, 3)
             assert torch.allclose(block.sum(-1), torch.ones(1), atol=1e-5), member
         exact = self._exact_member_marginals(fg, {})
         for member in h.members:
             assert torch.allclose(
-                out.probs["h"][..., h.column_of(member)], exact[member], atol=1e-5
+                out.probs[member], exact[member], atol=1e-5
             ), member
 
     def test_categorical_plate_partial_evidence(self):
@@ -723,14 +725,14 @@ def test_categorical_plate_partial_evidence(self):
         )
         exact = self._exact_member_marginals(fg, {"h1": 2, "g1": 1})
         h, g = fg.variables["h"], fg.variables["g"]
-        assert torch.allclose(out.probs["h"][..., h.column_of("h2")], exact["h2"], atol=1e-5)
+        assert torch.allclose(out.probs["h2"], exact["h2"], atol=1e-5)
         assert torch.allclose(
-            _state_marginal(g.member("g2"), out.probs["g"][..., g.column_of("g2")]),
+            _state_marginal(g.member("g2"), out.probs["g2"]),
             exact["g2"], atol=1e-5,
         )
         # observed members echo their evidence
         assert torch.allclose(
-            out.probs["h"][..., h.column_of("h1")], torch.tensor([[0.0, 0.0, 1.0]])
+            out.probs["h1"], torch.tensor([[0.0, 0.0, 1.0]])
         )
 
     def test_categorical_plate_trains(self):
@@ -745,7 +747,7 @@ def test_categorical_plate_trains(self):
         def loss_fn():
             out = eng.query(query=["h"], evidence=evidence)
             return torch.nn.functional.nll_loss(
-                out.probs["h"][..., h.column_of("h2")].log(), target
+                out.probs["h2"].log(), target
             )
 
         first = loss_fn().item()
@@ -756,4 +758,4 @@ def loss_fn():
             opt.step()
         assert loss.item() < first * 0.5
         out = eng.query(query=["h"], evidence=evidence)
-        assert out.probs["h"][0, h.column_of("h2")][2].item() > 0.9
+        assert out.probs["h2"][0][2].item() > 0.9