Added set_mean for part. wrapped normal dist

pyrecest/distributions/cart_prod/partially_wrapped_normal_distribution.py

@@ -108,9 +108,21 @@ def mode(self):
         """
         return self.mu
 
+    def set_mean(self, new_mean):
+        """
+        Return a copy of this distribution with the location parameter shifted to ``new_mean``.
+
+        For bounded dimensions, the mean is wrapped into [0, 2*pi) to stay on the manifold.
+        """
+        new_dist = copy.deepcopy(self)
+        wrapped_mean = where(
+            arange(new_mean.shape[0]) < self.bound_dim, mod(new_mean, 2 * pi), new_mean
+        )
+        new_dist.mu = wrapped_mean
+        return new_dist
+
     def set_mode(self, new_mode):
-        self.mu = copy.copy(new_mode)
-        return self
+        return self.set_mean(new_mode)
 
     def hybrid_moment(self):
         """