Matin 63 markov blanket normal form

pyproject.toml

@@ -25,3 +25,4 @@ cmake = ">=4.1.2,<5"
 pytest = ">=9.0.2,<10"
 numpy = ">=2.3.5,<3"
 basedpyright = ">=1.36.1,<2"
+ipython = ">=9.8.0,<10"

pyrightconfig.json

@@ -0,0 +1,4 @@
+{
+    "venvPath": ".pixi/envs",
+    "venv": "default"
+}

src/gibbsflux/conjugacy/normal_form.py

@@ -0,0 +1,21 @@
+from gibbsflux.factors.syntax import terms as t
+
+type KeyType = tuple[int | str | KeyType, ...]
+
+def key(term: t.Term, blanket_var: str) -> KeyType: # pyright: ignore
+    match term:
+        case t.Var(_, var):
+            return (0,) if var == blanket_var else (1, var)
+        case _ if t.is_base_factor_term(term):
+            return (2, *t.arg_positions(term, blanket_var), type(term).__name__) # pyright: ignore
+        case t.Mix(arg, branches):
+            return (3, int(arg != blanket_var),
+                    *(key(branch, blanket_var) for branch in branches))
+
+
+def normalize(markov_blanket: t.Program, blanket_var: str):
+    # precondition: markov_blanket is well-typed and actually encodes a Markov
+    # blanket factor graph
+    return t.Program(sorted(markov_blanket.terms,
+                            key=lambda term: key(term, blanket_var)))
+

src/gibbsflux/factors/syntax/terms.py

@@ -84,14 +84,20 @@ def __eq__(self, other):
 
 ## helpers ##
 
+def arg_positions(term: BaseFactor, var: str) -> tuple[int, ...]:
+    return tuple(i for (i, arg) in enumerate(astuple(term))\
+                 if isinstance(arg, str) and arg == var)
+
+def is_base_factor_term(term: Term) -> bool:
+    return any(isinstance(term, type) for type in BaseFactor.__args__)
+
 def free_vars(term: Term) -> set[str]:
     match term:
         case Var(_, name):
             return {name}
-        case Normal() | Dirichlet() | Categorical():
+        case _ if is_base_factor_term(term):
             return set(filter(lambda arg: isinstance(arg, str), astuple(term)))
         case Mix(arg, branches):
             argvars = {arg} if isinstance(arg, str) else set()
             return argvars.union(*map(free_vars, branches))
 
-