Skip to content

Allow copy and deepcopy of PYMC models #7492

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 9 commits into from
Oct 3, 2024
Merged
5 changes: 3 additions & 2 deletions pymc/model/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -1571,16 +1571,17 @@
def __contains__(self, key):
return key in self.named_vars or self.name_for(key) in self.named_vars

def __copy__(self):
return self.copy()

Check warning on line 1575 in pymc/model/core.py

View check run for this annotation

Codecov / codecov/patch

pymc/model/core.py#L1574-L1575

Added lines #L1574 - L1575 were not covered by tests

def __deepcopy__(self, _):
return self.copy()

Check warning on line 1578 in pymc/model/core.py

View check run for this annotation

Codecov / codecov/patch

pymc/model/core.py#L1578

Added line #L1578 was not covered by tests

def copy(self):
"""
Clone a pymc model by overiding the python copy method using the clone_model method from fgraph.
Constants are not cloned and if guassian process variables are detected then a warning will be triggered.
Clone the model

To access variables in the cloned model use `cloned_model["var_name"]`.

Examples
--------
Expand Down
42 changes: 18 additions & 24 deletions tests/model/test_core.py
Original file line number Diff line number Diff line change
Expand Up @@ -1765,17 +1765,13 @@ def test_graphviz_call_function(self, var_names, filenames) -> None:


class TestModelCopy:
@staticmethod
def simple_model() -> pm.Model:
@pytest.mark.parametrize("copy_method", (copy.copy, copy.deepcopy))
def test_copy_model(self, copy_method) -> None:
with pm.Model() as simple_model:
error = pm.HalfNormal("error", 0.5)
alpha = pm.Normal("alpha", 0, 1)
pm.Normal("y", alpha, error)
return simple_model

@pytest.mark.parametrize("copy_method", (copy.copy, copy.deepcopy))
def test_copy_model(self, copy_method) -> None:
simple_model = self.simple_model()
copy_simple_model = copy_method(simple_model)

with simple_model:
Expand All @@ -1786,15 +1782,24 @@ def test_copy_model(self, copy_method) -> None:
samples=1, random_seed=42
)

simple_model_prior_predictive_mean = simple_model_prior_predictive["prior"]["y"].mean(
("chain", "draw")
)
copy_simple_model_prior_predictive_mean = copy_simple_model_prior_predictive["prior"][
simple_model_prior_predictive_val = simple_model_prior_predictive["prior"]["y"].values
copy_simple_model_prior_predictive_val = copy_simple_model_prior_predictive["prior"][
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just compare directly, no need to assign to separate variables, that are almost as verbose as the way they are accessed

"y"
].mean(("chain", "draw"))
].values

assert np.isclose(
simple_model_prior_predictive_mean, copy_simple_model_prior_predictive_mean
assert simple_model_prior_predictive_val == copy_simple_model_prior_predictive_val

with copy_simple_model:
z = pm.Deterministic("z", copy_simple_model["alpha"] + 1)
copy_simple_model_prior_predictive = pm.sample_prior_predictive(
samples=1, random_seed=42
)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You can do this above, and call sample_prior_predictive only once for the copy_simple_model


assert "z" in copy_simple_model.named_vars
assert "z" not in simple_model.named_vars
assert (
copy_simple_model_prior_predictive["prior"]["z"].values
== 1 + copy_simple_model_prior_predictive["prior"]["alpha"].values
)

@pytest.mark.parametrize("copy_method", (copy.copy, copy.deepcopy))
Expand All @@ -1811,14 +1816,3 @@ def test_guassian_process_copy_failure(self, copy_method) -> None:
match="Detected variables likely created by GP objects. Further use of these old GP objects should be avoided as it may reintroduce variables from the old model. See issue: https://github.com/pymc-devs/pymc/issues/6883",
):
copy_method(gaussian_process_model)

@pytest.mark.parametrize("copy_method", (copy.copy, copy.deepcopy))
def test_adding_deterministics_to_clone(self, copy_method) -> None:
simple_model = self.simple_model()
clone_model = copy_method(simple_model)

with clone_model:
z = pm.Deterministic("z", clone_model["alpha"] + 1)

assert "z" in clone_model.named_vars
assert "z" not in simple_model.named_vars
Loading