Allow copy and deepcopy of PYMC models

pymc/model/core.py

@@ -1571,6 +1571,62 @@
     def __contains__(self, key):
         return key in self.named_vars or self.name_for(key) in self.named_vars
 
+    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.
+
+        Examples
+        --------
+        .. code-block:: python
+
+            import pymc as pm
+            import copy
+
+            with pm.Model() as m:
+                p = pm.Beta("p", 1, 1)
+                x = pm.Bernoulli("x", p=p, shape=(3,))
+
+            clone_m = copy.copy(m)
+
+            # Access cloned variables by name
+            clone_x = clone_m["x"]
+
+            # z will be part of clone_m but not m
+            z = pm.Deterministic("z", clone_x + 1)
+        """
+        from pymc.model.fgraph import clone_model
+
+        return clone_model(self)
+
+    def __deepcopy__(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.
+
+        Examples
+        --------
+        .. code-block:: python
+
+            import pymc as pm
+            import copy
+
+            with pm.Model() as m:
+                p = pm.Beta("p", 1, 1)
+                x = pm.Bernoulli("x", p=p, shape=(3,))
+
+            clone_m = copy.deepcopy(m)
+
+            # Access cloned variables by name
+            clone_x = clone_m["x"]
+
+            # z will be part of clone_m but not m
+            z = pm.Deterministic("z", clone_x + 1)
+        """
+        from pymc.model.fgraph import clone_model
+
+        return clone_model(self)
+
     def replace_rvs_by_values(
         self,
         graphs: Sequence[TensorVariable],

pymc/model/fgraph.py

@@ -11,6 +11,8 @@
 #   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
+import warnings
+
 from copy import copy, deepcopy
 
 import pytensor
@@ -158,6 +160,14 @@ def fgraph_from_model(
             "Nested sub-models cannot be converted to fgraph. Convert the parent model instead"
         )
 
+    if any(
+        ("_rotated_" in var_name or "_hsgp_coeffs_" in var_name) for var_name in model.named_vars
+    ):
+        warnings.warn(
+            "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",
+            UserWarning,
+        )
+
     # Collect PyTensor variables
     rvs_to_values = model.rvs_to_values
     rvs = list(rvs_to_values.keys())

tests/model/test_core.py

@@ -11,6 +11,7 @@
 #   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
+import copy
 import pickle
 import threading
 import traceback
@@ -1761,3 +1762,62 @@ def test_graphviz_call_function(self, var_names, filenames) -> None:
                 figsize=None,
                 dpi=300,
             )
+
+
+class TestModelCopy:
+    @staticmethod
+    def simple_model() -> pm.Model:
+        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
+
+    @staticmethod
+    def gp_model() -> pm.Model:
+        with pm.Model() as gp_model:
+            ell = pm.Gamma("ell", alpha=2, beta=1)
+            cov = 2 * pm.gp.cov.ExpQuad(1, ell)
+            gp = pm.gp.Latent(cov_func=cov)
+            f = gp.prior("f", X=np.arange(10)[:, None])
+            pm.Normal("y", f * 2)
+        return gp_model
+
+    def test_copy_model(self) -> None:
+        simple_model = self.simple_model()
+        copy_simple_model = copy.copy(simple_model)
+        deepcopy_simple_model = copy.deepcopy(simple_model)
+
+        with simple_model:
+            simple_model_prior_predictive = pm.sample_prior_predictive(random_seed=42)
+
+        with copy_simple_model:
+            copy_simple_model_prior_predictive = pm.sample_prior_predictive(random_seed=42)
+
+        with deepcopy_simple_model:
+            deepcopy_simple_model_prior_predictive = pm.sample_prior_predictive(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"][
+            "y"
+        ].mean(("chain", "draw"))
+        deepcopy_simple_model_prior_predictive_mean = deepcopy_simple_model_prior_predictive[
+            "prior"
+        ]["y"].mean(("chain", "draw"))
+
+        assert np.isclose(
+            simple_model_prior_predictive_mean, copy_simple_model_prior_predictive_mean
+        )
+        assert np.isclose(
+            simple_model_prior_predictive_mean, deepcopy_simple_model_prior_predictive_mean
+        )
+
+    def test_guassian_process_copy_failure(self) -> None:
+        gaussian_process_model = self.gp_model()
+        with pytest.warns(UserWarning):
+            copy.copy(gaussian_process_model)
+
+        with pytest.warns(UserWarning):
+            copy.deepcopy(gaussian_process_model)