WIP Improve scoped models (#5607)

* use slashes to separate context

* add fix for nested prefix

* fix test

* refactor prefix

* revert precommit change

* fix windows issue

* fis smc test

* Update pymc/tests/test_model.py

Co-authored-by: Ricardo Vieira <[email protected]>

* Update pymc/tests/test_model.py

Co-authored-by: Ricardo Vieira <[email protected]>

* update tests

Co-authored-by: Ricardo Vieira <[email protected]>

Author: ferrine

pymc/model.py

@@ -473,7 +473,7 @@ def __init__(self, mean=0, sigma=1, name=''):
 
                 # 3) you can create variables with Var method
                 self.Var('v1', Normal.dist(mu=mean, sigma=sd))
-                # this will create variable named like '{prefix_}v1'
+                # this will create variable named like '{prefix/}v1'
                 # and assign attribute 'v1' to instance created
                 # variable can be accessed with self.v1 or self['v1']
 
@@ -515,6 +515,8 @@ def __init__(self, mean=0, sigma=1, name=''):
             CustomModel(mean=1, name='first')
             CustomModel(mean=2, name='second')
 
+        # variables inside both scopes will be named like `first/*`, `second/*`
+
     """
 
     if TYPE_CHECKING:
@@ -1455,14 +1457,18 @@ def add_random_variable(self, var, dims: Optional[Tuple[Union[str, None], ...]]
             setattr(self, self.name_of(var.name), var)
 
     @property
-    def prefix(self):
-        return f"{self.name}_" if self.name else ""
+    def prefix(self) -> str:
+        if self.isroot or not self.parent.prefix:
+            name = self.name
+        else:
+            name = f"{self.parent.prefix}/{self.name}"
+        return name.strip("/")
 
     def name_for(self, name):
         """Checks if name has prefix and adds if needed"""
         if self.prefix:
             if not name.startswith(self.prefix):
-                return f"{self.prefix}{name}"
+                return f"{self.prefix}/{name}"
             else:
                 return name
         else:
@@ -1472,8 +1478,8 @@ def name_of(self, name):
         """Checks if name has prefix and deletes if needed"""
         if not self.prefix or not name:
             return name
-        elif name.startswith(self.prefix):
-            return name[len(self.prefix) :]
+        elif name.startswith(self.prefix + "/"):
+            return name[len(self.prefix) + 1 :]
         else:
             return name
 

pymc/tests/test_data_container.py

@@ -403,8 +403,8 @@ def test_data_naming():
     with pm.Model("named_model") as model:
         x = pm.ConstantData("x", [1.0, 2.0, 3.0])
         y = pm.Normal("y")
-    assert y.name == "named_model_y"
-    assert x.name == "named_model_x"
+    assert y.name == "named_model/y"
+    assert x.name == "named_model/x"
 
 
 def test_get_data():

pymc/tests/test_model.py

@@ -95,20 +95,20 @@ def test_context_passes_vars_to_parent_model(self):
                 usermodel2.register_rv(pm.Normal.dist(), "v3")
                 pm.Normal("v4")
                 # this variable is created in parent model too
-        assert "another_v2" in model.named_vars
-        assert "another_v3" in model.named_vars
-        assert "another_v3" in usermodel2.named_vars
-        assert "another_v4" in model.named_vars
-        assert "another_v4" in usermodel2.named_vars
+        assert "another/v2" in model.named_vars
+        assert "another/v3" in model.named_vars
+        assert "another/v3" in usermodel2.named_vars
+        assert "another/v4" in model.named_vars
+        assert "another/v4" in usermodel2.named_vars
         assert hasattr(usermodel2, "v3")
         assert hasattr(usermodel2, "v2")
         assert hasattr(usermodel2, "v4")
         # When you create a class based model you should follow some rules
         with model:
             m = NewModel("one_more")
-        assert m.d is model["one_more_d"]
-        assert m["d"] is model["one_more_d"]
-        assert m["one_more_d"] is model["one_more_d"]
+        assert m.d is model["one_more/d"]
+        assert m["d"] is model["one_more/d"]
+        assert m["one_more/d"] is model["one_more/d"]
 
 
 class TestNested:
@@ -124,8 +124,8 @@ def test_nest_context_works(self):
     def test_named_context(self):
         with pm.Model() as m:
             NewModel(name="new")
-        assert "new_v1" in m.named_vars
-        assert "new_v2" in m.named_vars
+        assert "new/v1" in m.named_vars
+        assert "new/v2" in m.named_vars
 
     def test_docstring_example1(self):
         usage1 = DocstringModel()
@@ -138,10 +138,10 @@ def test_docstring_example1(self):
     def test_docstring_example2(self):
         with pm.Model() as model:
             DocstringModel(name="prefix")
-        assert "prefix_v1" in model.named_vars
-        assert "prefix_v2" in model.named_vars
-        assert "prefix_v3" in model.named_vars
-        assert "prefix_v3_sq" in model.named_vars
+        assert "prefix/v1" in model.named_vars
+        assert "prefix/v2" in model.named_vars
+        assert "prefix/v3" in model.named_vars
+        assert "prefix/v3_sq" in model.named_vars
         assert len(model.potentials), 1
 
     def test_duplicates_detection(self):
@@ -156,6 +156,20 @@ def test_model_root(self):
             with pm.Model() as sub:
                 assert model is sub.root
 
+    def test_nested_named_model_repeated(self):
+        with pm.Model("sub") as model:
+            b = pm.Normal("var")
+            with pm.Model("sub"):
+                b = pm.Normal("var")
+        assert {"sub/var", "sub/sub/var"} == set(model.named_vars.keys())
+
+    def test_nested_named_model(self):
+        with pm.Model("sub1") as model:
+            b = pm.Normal("var")
+            with pm.Model("sub2"):
+                b = pm.Normal("var")
+        assert {"sub1/var", "sub1/sub2/var"} == set(model.named_vars.keys())
+
 
 class TestObserved:
     def test_observed_rv_fail(self):
@@ -658,14 +672,14 @@ def test_datalogpt_multiple_shapes():
 
 
 def test_nested_model_coords():
-    COORDS = {"dim": range(10)}
-    with pm.Model(name="m1", coords=COORDS) as m1:
-        a = pm.Normal("a")
-        with pm.Model(name="m2") as m2:
-            b = pm.Normal("b")
-            c = pm.HalfNormal("c")
-            d = pm.Normal("d", b, c, dims="dim")
-        e = pm.Normal("e", a + d, dims="dim")
+    with pm.Model(name="m1", coords=dict(dim1=range(2))) as m1:
+        a = pm.Normal("a", dims="dim1")
+        with pm.Model(name="m2", coords=dict(dim2=range(4))) as m2:
+            b = pm.Normal("b", dims="dim1")
+            m1.add_coord("dim3", range(4))
+            c = pm.HalfNormal("c", dims="dim3")
+            d = pm.Normal("d", b, c, dims="dim2")
+        e = pm.Normal("e", a[None] + d[:, None], dims=("dim2", "dim1"))
     assert m1.coords is m2.coords
     assert m1.dim_lengths is m2.dim_lengths
     assert set(m2.RV_dims) < set(m1.RV_dims)

pymc/tests/test_smc.py

@@ -530,9 +530,9 @@ def test_named_model(self):
             s = pm.Simulator("s", self.normal_sim, a, b, observed=self.data)
 
             trace = pm.sample_smc(draws=10, chains=2, return_inferencedata=False)
-            assert f"{name}_a" in trace.varnames
-            assert f"{name}_b" in trace.varnames
-            assert f"{name}_b_log__" in trace.varnames
+            assert f"{name}/a" in trace.varnames
+            assert f"{name}/b" in trace.varnames
+            assert f"{name}/b_log__" in trace.varnames
 
 
 class TestMHKernel(SeededTest):