CUQIpy 1.0.0.post0.dev229__py3-none-any.whl → 1.0.0.post0.dev305__py3-none-any.whl

{CUQIpy-1.0.0.post0.dev229.dist-info → CUQIpy-1.0.0.post0.dev305.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: CUQIpy
-Version: 1.0.0.post0.dev229
+Version: 1.0.0.post0.dev305
 Summary: Computational Uncertainty Quantification for Inverse problems in Python
 Maintainer-email: "Nicolai A. B. Riis" <nabr@dtu.dk>, "Jakob S. Jørgensen" <jakj@dtu.dk>, "Amal M. Alghamdi" <amaal@dtu.dk>, Chao Zhang <chaz@dtu.dk>
 License: Apache License

{CUQIpy-1.0.0.post0.dev229.dist-info → CUQIpy-1.0.0.post0.dev305.dist-info}/RECORD

@@ -1,6 +1,6 @@
 cuqi/__init__.py,sha256=LsGilhl-hBLEn6Glt8S_l0OJzAA1sKit_rui8h-D-p0,488
 cuqi/_messages.py,sha256=fzEBrZT2kbmfecBBPm7spVu7yHdxGARQB4QzXhJbCJ0,415
-cuqi/_version.py,sha256=qSCcU546LTLb2CnsgdAFPFHEfSeVUxNMie_SdR9XFeY,510
+cuqi/_version.py,sha256=Lro7539e-VFksYWNf7xFPv5bxAlSPs5DaTne7ms-G4g,510
 cuqi/config.py,sha256=wcYvz19wkeKW2EKCGIKJiTpWt5kdaxyt4imyRkvtTRA,526
 cuqi/diagnostics.py,sha256=5OrbJeqpynqRXOe5MtOKKhe7EAVdOEpHIqHnlMW9G_c,3029
 cuqi/array/__init__.py,sha256=-EeiaiWGNsE3twRS4dD814BIlfxEsNkTCZUc5gjOXb0,30
@@ -33,14 +33,14 @@ cuqi/distribution/_posterior.py,sha256=zAfL0GECxekZ2lBt1W6_LN0U_xskMwK4VNce5xAF7
 cuqi/distribution/_uniform.py,sha256=7xJmCZH_LPhuGkwEDGh-_CTtzcWKrXMOxtTJUFb7Ydo,1607
 cuqi/experimental/__init__.py,sha256=vhZvyMX6rl8Y0haqCzGLPz6PSUKyu75XMQbeDHqTTrw,83
 cuqi/experimental/mcmc/__init__.py,sha256=UqoyPWNQt4ZGIgc9Buhl5gf3toAxLjXLyQ7DieDQlRw,384
-cuqi/experimental/mcmc/_cwmh.py,sha256=yRlTk5a1QYfH3JyCecfOOTeDf-4-tmJ3Tl2Bc3pyp1Y,7336
-cuqi/experimental/mcmc/_hmc.py,sha256=qqAyoAajLE_JenYMgAbD3tknuEf75AJu-ufF69GKGk4,19384
-cuqi/experimental/mcmc/_langevin_algorithm.py,sha256=MX48u3GYgCckB6Q5h5kXr_qdIaLQH2toOG5u29OY7gk,8245
-cuqi/experimental/mcmc/_laplace_approximation.py,sha256=7reeOnDY77WnOwqYls5WStftHgylwCNVodudRroApF0,5812
-cuqi/experimental/mcmc/_mh.py,sha256=aIV1Ntq0EAq3QJ1_X-DbP7eDAL-d_Or7d3RUO-R48I4,3090
-cuqi/experimental/mcmc/_pcn.py,sha256=m7pR266uUJQociOe_CpUUlKHkfm8g--JfRWaQA2IKis,4364
-cuqi/experimental/mcmc/_rto.py,sha256=jSPznr34XPfWM6LysWIiN4hE-vtyti3cHyvzy9ruykg,11349
-cuqi/experimental/mcmc/_sampler.py,sha256=_5Uo2F-Mta46w3lo7WBVNwvTLYhES_BzMTJrKxA00c8,14861
+cuqi/experimental/mcmc/_cwmh.py,sha256=dbvmy6Fyr_xszbO3YmYcRsDaRRtQfPimLo6rbp0II6M,6898
+cuqi/experimental/mcmc/_hmc.py,sha256=0sZMHtnNFGGtQdzpx-cgqA0xyfvGy7r4K62RH3AQNa4,19285
+cuqi/experimental/mcmc/_langevin_algorithm.py,sha256=n6WRQooKuUDjmqF-CtpcSNFDvaHCgLKhWxX-hi7h_ZA,8224
+cuqi/experimental/mcmc/_laplace_approximation.py,sha256=xrlQbRaoJ8Rx1iWvSqrxX40OsirmsRuEO3R4ewuZdGM,5702
+cuqi/experimental/mcmc/_mh.py,sha256=DUoZJ1cFnHkupmTV6-GIbINLXhG4stps44RVlkuGQ8s,2625
+cuqi/experimental/mcmc/_pcn.py,sha256=T4T32mfoii3k6Jfz0qxPQbwdh6wdVOxttiEP7NWaZzg,3386
+cuqi/experimental/mcmc/_rto.py,sha256=49RopzzbfRp6c_7WTFgF7e602hb4YHVuls9YWRfCmGk,10081
+cuqi/experimental/mcmc/_sampler.py,sha256=g9c1ds5PoXr4LDYF3h_1wFI1M6m2nYp1poVu0RHrwh8,19720
 cuqi/geometry/__init__.py,sha256=Tz1WGzZBY-QGH3c0GiyKm9XHN8MGGcnU6TUHLZkzB3o,842
 cuqi/geometry/_geometry.py,sha256=WYFC-4_VBTW73b2ldsnfGYKvdSiCE8plr89xTSmkadg,46804
 cuqi/implicitprior/__init__.py,sha256=ZRZ9fgxgEl5n0A9F7WCl1_jid-GUiC8ZLkyTmGQmFlY,100
@@ -77,8 +77,8 @@ cuqi/testproblem/_testproblem.py,sha256=x769LwwRdJdzIiZkcQUGb_5-vynNTNALXWKato7s
 cuqi/utilities/__init__.py,sha256=T4tLsC215MknBCsw_C0Qeeg_ox26aDUrCA5hbWvNQkU,387
 cuqi/utilities/_get_python_variable_name.py,sha256=QwlBVj2koJRA8s8pWd554p7-ElcI7HUwY32HknaR92E,1827
 cuqi/utilities/_utilities.py,sha256=MWAqV6L5btMpWwlUzrZYuV2VeSpfTbOaLRMRkuw2WIA,8509
-CUQIpy-1.0.0.post0.dev229.dist-info/LICENSE,sha256=kJWRPrtRoQoZGXyyvu50Uc91X6_0XRaVfT0YZssicys,10799
-CUQIpy-1.0.0.post0.dev229.dist-info/METADATA,sha256=-LQQOopcYpRJTp_ZhyH2H97JxI2jLfgsIS1AzYs1FNU,18393
-CUQIpy-1.0.0.post0.dev229.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-CUQIpy-1.0.0.post0.dev229.dist-info/top_level.txt,sha256=AgmgMc6TKfPPqbjV0kvAoCBN334i_Lwwojc7HE3ZwD0,5
-CUQIpy-1.0.0.post0.dev229.dist-info/RECORD,,
+CUQIpy-1.0.0.post0.dev305.dist-info/LICENSE,sha256=kJWRPrtRoQoZGXyyvu50Uc91X6_0XRaVfT0YZssicys,10799
+CUQIpy-1.0.0.post0.dev305.dist-info/METADATA,sha256=dqQXHzSfLCCBGnno4jP-AkVQDWa0u_RJR-QGFUXZKk8,18393
+CUQIpy-1.0.0.post0.dev305.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+CUQIpy-1.0.0.post0.dev305.dist-info/top_level.txt,sha256=AgmgMc6TKfPPqbjV0kvAoCBN334i_Lwwojc7HE3ZwD0,5
+CUQIpy-1.0.0.post0.dev305.dist-info/RECORD,,

cuqi/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2024-05-21T12:48:50+0200",
+ "date": "2024-05-27T07:34:21+0200",
  "dirty": false,
  "error": null,
- "full-revisionid": "eb9519734f2558f66772895e41f2cd0c3cd61767",
- "version": "1.0.0.post0.dev229"
+ "full-revisionid": "1274ccfff1cb9663b547dd2aeb9c86c72dccb2ae",
+ "version": "1.0.0.post0.dev305"
 }
 '''  # END VERSION_JSON
 

cuqi/experimental/mcmc/_cwmh.py

@@ -67,16 +67,21 @@ class CWMHNew(ProposalBasedSamplerNew):
         samples = sampler.sample(2000).get_samples()
 
     """
-    def __init__(self, target: cuqi.density.Density, proposal=None, scale=1,
+
+    _STATE_KEYS = ProposalBasedSamplerNew._STATE_KEYS.union(['_scale_temp'])
+
+    def __init__(self, target:cuqi.density.Density=None, proposal=None, scale=1,
                  initial_point=None, **kwargs):
         super().__init__(target, proposal=proposal, scale=scale,
                          initial_point=initial_point, **kwargs)
+        
+    def _initialize(self):
+        if isinstance(self.scale, Number):
+            self.scale = np.ones(self.dim)*self.scale
+        self._acc = [np.ones((self.dim))] # Overwrite acc from ProposalBasedSamplerNew with list of arrays
 
-        # set initial scale
-        self.scale = scale
-
-        # set initial acceptance rate
-        self._acc = [np.ones((self.dim))]
+        # Handling of temporary scale parameter due to possible bug in old CWMH
+        self._scale_temp = self.scale.copy()
 
     @property
     def scale(self):
@@ -86,53 +91,35 @@ class CWMHNew(ProposalBasedSamplerNew):
     @scale.setter
     def scale(self, value):
         """ Set the scale parameter. """
-        if isinstance(value, Number):
-            self._scale = np.ones(self.dim)*value
-        elif isinstance(value, np.ndarray):
-            self._scale = value
-        self._scale_temp = self._scale.copy()
+        if self._is_initialized and isinstance(value, Number):
+            value = np.ones(self.dim)*value
+        self._scale = value
 
     def validate_target(self):
         if not isinstance(self.target, cuqi.density.Density):
             raise ValueError(
                 "Target should be an instance of "+\
                 f"{cuqi.density.Density.__class__.__name__}")
-
-    @ProposalBasedSamplerNew.proposal.setter
-    # TODO. Check if we can refactor this.
-    # We can work with a validate_proposal method instead?
-    def proposal(self, value):
-        fail_msg = "Proposal should be either None, "+\
-            f"{cuqi.distribution.Distribution.__class__.__name__} "+\
-            "conditioned only on 'location' and 'scale', lambda function, "+\
-            f"or {cuqi.distribution.Normal.__class__.__name__} conditioned "+\
-            "only on 'mean' and 'std'"
-
-        if value is None:
+        
+    def validate_proposal(self):
+        if not isinstance(self.proposal, cuqi.distribution.Distribution):
+            raise ValueError("Proposal must be a cuqi.distribution.Distribution object")
+        if not self.proposal.is_symmetric:
+            raise ValueError("Proposal must be symmetric")
+        
+    @property
+    def proposal(self):
+        if self._proposal is None:
             self._proposal = cuqi.distribution.Normal(
                 mean=lambda location: location,
                 std=lambda scale: scale,
                 geometry=self.dim,
             )
-
-        elif isinstance(value, cuqi.distribution.Distribution) and sorted(
-            value.get_conditioning_variables()
-        ) == ["location", "scale"]:
-            self._proposal = value
-
-        elif isinstance(value, cuqi.distribution.Normal) and sorted(
-            value.get_conditioning_variables()
-        ) == ["mean", "std"]:
-            self._proposal = value(
-                mean=lambda location: location, std=lambda scale: scale
-            )
-
-        elif not isinstance(value, cuqi.distribution.Distribution) and callable(
-            value):
-            self._proposal = value
-
-        else:
-            raise ValueError(fail_msg)
+        return self._proposal
+    
+    @proposal.setter
+    def proposal(self, value):
+        self._proposal = value
 
     def step(self):
         # Initialize x_t which is used to store the current CWMH sample

cuqi/experimental/mcmc/_hmc.py

@@ -95,7 +95,7 @@ class NUTSNew(SamplerNew):
                                                     'epsilon_list',
                                                     'epsilon_bar_list'})
 
-    def __init__(self, target, initial_point=None, max_depth=15,
+    def __init__(self, target=None, initial_point=None, max_depth=15,
                  step_size=None, opt_acc_rate=0.6, **kwargs):
         super().__init__(target, initial_point=initial_point, **kwargs)
 
@@ -103,9 +103,6 @@ class NUTSNew(SamplerNew):
         self.max_depth = max_depth
         self.step_size = step_size
         self.opt_acc_rate = opt_acc_rate
-        
-        # Set current point 
-        self.current_point = self.initial_point
 
         # Initialize epsilon and epsilon_bar
         # epsilon is the step size used in the current iteration
@@ -115,12 +112,37 @@ class NUTSNew(SamplerNew):
         self._epsilon_bar = None
         self._H_bar = None
 
+        # Extra parameters for tuning
+        self._n_alpha = None
+        self._alpha = None
+
+
+    def _initialize(self):
+
         # Arrays to store acceptance rate
-        self._acc = [None]
+        self._acc = [None] # Overwrites acc from SamplerNew. TODO. Check if this is necessary
+
+        self._alpha = 0 # check if meaningful value
+        self._n_alpha = 0 # check if meaningful value
+
+        self.current_target_logd, self.current_target_grad = self._nuts_target(self.current_point)
+
+        # parameters dual averaging
+        if self.step_size is None:
+            self._epsilon = self._FindGoodEpsilon()
+        else:
+            self._epsilon = self.step_size
+        self._epsilon_bar = "unset"
+
+        # Parameter mu, does not change during the run
+        self._mu = np.log(10*self._epsilon)
+
+        self._H_bar = 0
 
         # NUTS run diagnostic:
         # number of tree nodes created each NUTS iteration
         self._num_tree_node = 0
+
         # Create lists to store NUTS run diagnostics
         self._create_run_diagnostic_attributes()
 
@@ -176,9 +198,9 @@ class NUTSNew(SamplerNew):
         except:
             raise ValueError('Target must have logd and gradient methods.')
 
-    def reset(self):
+    def reinitialize(self):
         # Call the parent reset method
-        super().reset()
+        super().reinitialize()
         # Reset NUTS run diagnostic attributes
         self._reset_run_diagnostic_attributes()
 
@@ -275,10 +297,6 @@ class NUTSNew(SamplerNew):
             np.exp(eta*np.log(self._epsilon) +(1-eta)*np.log(self._epsilon_bar))
 
     def _pre_warmup(self):
-        super()._pre_warmup()
-
-        self.current_target_logd, self.current_target_grad =\
-            self._nuts_target(self.current_point)
 
         # Set up tuning parameters (only first time tuning is called)
         # Note:
@@ -289,32 +307,14 @@ class NUTSNew(SamplerNew):
         #  Parameters that does not change during the run
         #    self._mu
 
-        if self._epsilon is None:
-            # parameters dual averaging
-            self._epsilon = self._FindGoodEpsilon()
-            # Parameter mu, does not change during the run
-            self._mu = np.log(10*self._epsilon)
-
-        if self._epsilon_bar is None: # Initial value of epsilon_bar
+        if self._epsilon_bar == "unset": # Initial value of epsilon_bar for tuning
             self._epsilon_bar = 1
 
-        if self._H_bar is None: # Initial value of H_bar
-            self._H_bar = 0
-
     def _pre_sample(self):
-        super()._pre_sample()
 
-        self.current_target_logd, self.current_target_grad =\
-            self._nuts_target(self.current_point)
-        
-        # Set up epsilon and epsilon_bar if not set
-        if self._epsilon is None:
-            if self.step_size is None:
-                step_size = self._FindGoodEpsilon()
-            else:
-                step_size = self.step_size
-            self._epsilon = step_size
-            self._epsilon_bar = step_size
+        if self._epsilon_bar == "unset": # Initial value of epsilon_bar for sampling
+            self._epsilon_bar = self._epsilon
+            
 
     #=========================================================================
     def _nuts_target(self, x): # returns logposterior tuple evaluation-gradient
@@ -467,4 +467,4 @@ class NUTSNew(SamplerNew):
         # Store the step size used in iteration k
         self.epsilon_list.append(eps)
         # Store the step size suggestion during adaptation in iteration k
-        self.epsilon_bar_list.append(eps_bar)
+        self.epsilon_bar_list.append(eps_bar)

cuqi/experimental/mcmc/_langevin_algorithm.py

@@ -63,15 +63,16 @@ class ULANew(SamplerNew): # Refactor to Proposal-based sampler?
 
     _STATE_KEYS = SamplerNew._STATE_KEYS.union({'current_target_logd', 'scale', 'current_target_grad'})
 
-    def __init__(self, target, scale=1.0, **kwargs):
+    def __init__(self, target=None, scale=1.0, **kwargs):
 
         super().__init__(target, **kwargs)
 
-        self.scale = scale
-        self.current_point = self.initial_point
+        self.initial_scale = scale
+
+    def _initialize(self):
+        self.scale = self.initial_scale
         self.current_target_logd = self.target.logd(self.current_point)
         self.current_target_grad = self.target.gradient(self.current_point)
-        self._acc = [1] # TODO. Check if we need this
 
     def validate_target(self):
         try:

cuqi/experimental/mcmc/_laplace_approximation.py

@@ -53,17 +53,13 @@ class UGLANew(SamplerNew):
 
         super().__init__(target=target, initial_point=initial_point, **kwargs)
 
-        if initial_point is None: #TODO: Replace later with a getter
-            self.initial_point = np.zeros(self.dim)
-            self._samples = [self.initial_point]
-
-        self.current_point = self.initial_point
-        self._acc = [1] # TODO. Check if we need this
-
         # Parameters
         self.maxit = maxit
         self.tol = tol
         self.beta = beta
+    
+    def _initialize(self):
+        self._precompute()
 
     @property
     def prior(self):
@@ -81,8 +77,8 @@ class UGLANew(SamplerNew):
     def data(self):
         return self.target.data
 
-    def _pre_warmup(self):
-        super()._pre_warmup()
+    def _precompute(self):
+
         D = self.prior._diff_op
         n = D.shape[0]
 
@@ -121,9 +117,6 @@ class UGLANew(SamplerNew):
             return out
         self.M = M
 
-    def _pre_sample(self):
-        self._pre_warmup()
-
     def step(self):
         # Update Laplace approximation
         self._L2 = self.Lk_fun(self.current_point)
@@ -157,3 +150,8 @@ class UGLANew(SamplerNew):
         # Check that prior is LMRF
         if not isinstance(self.prior, cuqi.distribution.LMRF):
             raise ValueError('Unadjusted Gaussian Laplace approximation (UGLA) requires LMRF prior')
+        
+    @property
+    def _default_initial_point(self):
+        """ Get the default initial point for the sampler. Defaults to an array of zeros. """
+        return np.zeros(self.dim)

cuqi/experimental/mcmc/_mh.py

@@ -1,7 +1,6 @@
 import numpy as np
 import cuqi
 from cuqi.experimental.mcmc import ProposalBasedSamplerNew
-from cuqi.array import CUQIarray
 
 
 class MHNew(ProposalBasedSamplerNew):
@@ -27,25 +26,19 @@ class MHNew(ProposalBasedSamplerNew):
 
     def __init__(self, target, proposal=None, scale=1, **kwargs):
         super().__init__(target, proposal=proposal, scale=scale, **kwargs)
+
+    def _initialize(self):
         # Due to a bug? in old MH, we must keep track of this extra variable to match behavior.
-        self._scale_temp = self.scale 
+        self._scale_temp = self.scale
 
     def validate_target(self):
         pass # All targets are valid
 
-    @ProposalBasedSamplerNew.proposal.setter  # TODO. Check if we can refactor this. We can work with a validate_proposal method instead?
-    def proposal(self, value):
-        fail_msg = "Proposal should be either None, symmetric cuqi.distribution.Distribution or a lambda function."
-
-        if value is None:
-            self._proposal = cuqi.distribution.Gaussian(np.zeros(self.dim), 1)
-        elif not isinstance(value, cuqi.distribution.Distribution) and callable(value):
-            raise NotImplementedError(fail_msg)
-        elif isinstance(value, cuqi.distribution.Distribution) and value.is_symmetric:
-            self._proposal = value
-        else:
-            raise ValueError(fail_msg)
-        self._proposal.geometry = self.target.geometry
+    def validate_proposal(self):
+        if not isinstance(self.proposal, cuqi.distribution.Distribution):
+            raise ValueError("Proposal must be a cuqi.distribution.Distribution object")
+        if not self.proposal.is_symmetric:
+            raise ValueError("Proposal must be symmetric")
 
     def step(self):
         # propose state

cuqi/experimental/mcmc/_pcn.py

@@ -7,29 +7,25 @@ class PCNNew(SamplerNew):  # Refactor to Proposal-based sampler?
 
     _STATE_KEYS = SamplerNew._STATE_KEYS.union({'scale', 'current_likelihood_logd', 'lambd'})
 
-    def __init__(self, target, scale=1.0, **kwargs):
+    def __init__(self, target=None, scale=1.0, **kwargs):
 
         super().__init__(target, **kwargs)
+        self.initial_scale = scale
 
-        self.scale = scale
-        self.current_point = self.initial_point
+    def _initialize(self):
+        self.scale = self.initial_scale
         self.current_likelihood_logd = self._loglikelihood(self.current_point)
 
-        self._acc = [1] # TODO. Check if we need this
-
         # parameters used in the Robbins-Monro recursion for tuning the scale parameter
         # see details and reference in the tune method
         self.lambd = self.scale
         self.star_acc = 0.44 #TODO: 0.234 # target acceptance rate
 
     def validate_target(self):
-        try:
-            if isinstance(self.prior, (cuqi.distribution.Gaussian, cuqi.distribution.Normal)):
-                pass
-            else:
-                raise ValueError("The prior distribution of the target need to be Gaussian")
-        except AttributeError:
-            raise ValueError("The target need to have a prior distribution")
+        if not isinstance(self.target, cuqi.distribution.Posterior):
+            raise ValueError(f"To initialize an object of type {self.__class__}, 'target' need to be of type 'cuqi.distribution.Posterior'.")
+        if not isinstance(self.prior, (cuqi.distribution.Gaussian, cuqi.distribution.Normal)):
+            raise ValueError("The prior distribution of the target need to be Gaussian")
 
     def step(self):
         # propose state
@@ -55,30 +51,14 @@ class PCNNew(SamplerNew):  # Refactor to Proposal-based sampler?
 
     @property
     def prior(self):
-        if isinstance(self.target, cuqi.distribution.Posterior):
-            return self.target.prior
-        elif isinstance(self.target,tuple) and len(self.target)==2:
-            return self.target[1]
+        return self.target.prior
 
     @property
     def likelihood(self):
-        if isinstance(self.target, cuqi.distribution.Posterior):
-            return self.target.likelihood
-        elif isinstance(self.target,tuple) and len(self.target)==2:
-            return self.target[0]
-
-    @SamplerNew.target.setter 
-    def target(self, value):
-        if isinstance(value, cuqi.distribution.Posterior):
-            self._target = value
-            self._loglikelihood = lambda x : self.likelihood.logd(x)
-        elif isinstance(value,tuple) and len(value)==2 and \
-             (isinstance(value[0], cuqi.likelihood.Likelihood) or isinstance(value[0], cuqi.likelihood.UserDefinedLikelihood))  and \
-             isinstance(value[1], cuqi.distribution.Distribution):
-            self._target = value
-            self._loglikelihood = lambda x : self.likelihood.logd(x)
-        else:
-            raise ValueError(f"To initialize an object of type {self.__class__}, 'target' need to be of type 'cuqi.distribution.Posterior'.")
+        return self.target.likelihood
+        
+    def _loglikelihood(self, x):
+        return self.likelihood.logd(x)
 
     @property
     def dim(self): # TODO. Check if we need this. Implemented in base class
@@ -106,4 +86,4 @@ class PCNNew(SamplerNew):  # Refactor to Proposal-based sampler?
         self.lambd = np.exp(np.log(self.lambd) + zeta*(hat_acc-self.star_acc))
 
         # update scale parameter
-        self.scale = min(self.lambd, 1)
+        self.scale = min(self.lambd, 1)

cuqi/experimental/mcmc/_rto.py

@@ -5,7 +5,6 @@ import numpy as np
 import cuqi
 from cuqi.solver import CGLS, FISTA
 from cuqi.experimental.mcmc import SamplerNew
-from cuqi.array import CUQIarray
 
 
 class LinearRTONew(SamplerNew):
@@ -44,21 +43,17 @@ class LinearRTONew(SamplerNew):
         An example is shown in demos/demo31_callback.py.
         
     """
-    def __init__(self, target, initial_point=None, maxit=10, tol=1e-6, **kwargs):
+    def __init__(self, target=None, initial_point=None, maxit=10, tol=1e-6, **kwargs):
 
         super().__init__(target=target, initial_point=initial_point, **kwargs)
 
-        if initial_point is None: #TODO: Replace later with a getter
-            self.initial_point = np.zeros(self.dim)
-            self._samples = [self.initial_point]
-
-        self.current_point = self.initial_point
-        self._acc = [1] # TODO. Check if we need this
-
         # Other parameters
         self.maxit = maxit
         self.tol = tol
 
+    def _initialize(self):
+        self._precompute()
+
     @property
     def prior(self):
         return self.target.prior
@@ -81,41 +76,7 @@ class LinearRTONew(SamplerNew):
     @property
     def data(self):
         return self.target.data
-
-    @SamplerNew.target.setter
-    def target(self, value):
-        """ Set the target density. Runs validation of the target. """
-        # Accept tuple of inputs and construct posterior
-        if isinstance(value, tuple) and len(value) == 5:
-            # Structure (data, model, L_sqrtprec, P_mean, P_sqrtprec)
-            data = value[0]
-            model = value[1]
-            L_sqrtprec = value[2]
-            P_mean = value[3]
-            P_sqrtprec = value[4]
-
-            # If numpy matrix convert to CUQI model
-            if isinstance(model, np.ndarray) and len(model.shape) == 2:
-                model = cuqi.model.LinearModel(model)
-
-            # Check model input
-            if not isinstance(model, cuqi.model.LinearModel):
-                raise TypeError("Model needs to be cuqi.model.LinearModel or matrix")
-
-            # Likelihood
-            L = cuqi.distribution.Gaussian(model, sqrtprec=L_sqrtprec).to_likelihood(data)
-
-            # Prior TODO: allow multiple priors stacked
-            #if isinstance(P_mean, list) and isinstance(P_sqrtprec, list):
-            #    P = cuqi.distribution.JointGaussianSqrtPrec(P_mean, P_sqrtprec)
-            #else:
-            P = cuqi.distribution.Gaussian(P_mean, sqrtprec=P_sqrtprec)
-
-            # Construct posterior
-            value = cuqi.distribution.Posterior(L, P)
-        super(LinearRTONew, type(self)).target.fset(self, value)
-        self._precompute()
-
+    
     def _precompute(self):
         L1 = [likelihood.distribution.sqrtprec for likelihood in self.likelihoods]
         L2 = self.prior.sqrtprec
@@ -188,6 +149,11 @@ class LinearRTONew(SamplerNew):
 
         if not hasattr(self.prior, "sqrtprecTimesMean"):
             raise TypeError("Prior must contain a sqrtprecTimesMean attribute")
+    
+    @property
+    def _default_initial_point(self):
+        """ Get the default initial point for the sampler. Defaults to an array of zeros. """
+        return np.zeros(self.dim)
 
 class RegularizedLinearRTONew(LinearRTONew):
     """
@@ -231,15 +197,22 @@ class RegularizedLinearRTONew(LinearRTONew):
         self.stepsize = stepsize
         self.abstol = abstol   
         self.adaptive = adaptive
-        self.proximal = target.prior.proximal
-        self._stepsize = self._choose_stepsize()
         self.maxit = maxit
 
-    @LinearRTONew.target.setter
-    def target(self, value):
-        if not callable(value.prior.proximal):
-            raise TypeError("Projector needs to be callable")
-        return super(RegularizedLinearRTONew, type(self)).target.fset(self, value)
+    def _initialize(self):
+        super()._initialize()
+        self._stepsize = self._choose_stepsize()
+
+    @property
+    def proximal(self):
+        return self.target.prior.proximal
+    
+    def validate_target(self):
+        super().validate_target()
+        if not isinstance(self.target.prior, (cuqi.implicitprior.RegularizedGaussian, cuqi.implicitprior.RegularizedGMRF)):
+            raise TypeError("Prior needs to be RegularizedGaussian or RegularizedGMRF")
+        if not callable(self.proximal):
+            raise TypeError("Proximal needs to be callable")
 
     def _choose_stepsize(self):
         if isinstance(self.stepsize, str):

cuqi/experimental/mcmc/_sampler.py

@@ -20,7 +20,23 @@ class SamplerNew(ABC):
 
     Samples are stored in a list to allow for dynamic growth of the sample set. Returning samples is done by creating a new Samples object from the list of samples.
 
+    The sampler maintains sets of state and history keys, which are used for features like checkpointing and resuming sampling.
+
+    The state of the sampler represents all variables that are updated (replaced) in a Markov Monte Carlo step, e.g. the current point of the sampler.
+
+    The history of the sampler represents all variables that are updated (appended) in a Markov Monte Carlo step, e.g. the samples and acceptance rates.
+
+    Subclasses should ensure that any new variables that are updated in a Markov Monte Carlo step are added to the state or history keys.
+
+    Saving and loading checkpoints saves and loads the state of the sampler (not the history).
+
+    Batching samples via the batch_size parameter saves the sampler history to disk in batches of the specified size.
+
+    Any other attribute stored as part of the sampler (e.g. target, initial_point) is not supposed to be updated
+    during sampling and should not be part of the state or history.
+
     """
+
     _STATE_KEYS = {'current_point'}
     """ Set of keys for the state dictionary. """
 
@@ -29,6 +45,10 @@ class SamplerNew(ABC):
 
     def __init__(self, target: cuqi.density.Density, initial_point=None, callback=None):
         """ Initializer for abstract base class for all samplers.
+
+        Any subclassing samplers should simply store input parameters as part of the __init__ method. 
+
+        The actual initialization of the sampler should be done in the _initialize method.
         
         Parameters
         ----------
@@ -45,18 +65,37 @@ class SamplerNew(ABC):
         """
 
         self.target = target
+        self.initial_point = initial_point
         self.callback = callback
+        self._is_initialized = False
 
-        # Choose initial point if not given
-        if initial_point is None:
-            initial_point = np.ones(self.dim)
+    def initialize(self):
+        """ Initialize the sampler by setting and allocating the state and history before sampling starts. """
 
-        self.initial_point = initial_point
+        if self._is_initialized:
+            raise ValueError("Sampler is already initialized.")
         
-        self._samples = [initial_point] # Remove. See #324.
+        if self.target is None:
+            raise ValueError("Cannot initialize sampler without a target density.")
+        
+        # Default values
+        if self.initial_point is None:
+            self.initial_point = self._default_initial_point
+
+        # State variables
+        self.current_point = self.initial_point
+
+        # History variables
+        self._samples = []
+        self._acc = [ 1 ] # TODO. Check if we need to put 1 here.
+
+        self._initialize() # Subclass specific initialization
+
+        self._validate_initialization()
+
+        self._is_initialized = True
 
     # ------------ Abstract methods to be implemented by subclasses ------------
-    
     @abstractmethod
     def step(self):
         """ Perform one step of the sampler by transitioning the current point to a new point according to the sampler's transition kernel. """
@@ -72,23 +111,19 @@ class SamplerNew(ABC):
         """ Validate the target is compatible with the sampler. Called when the target is set. Should raise an error if the target is not compatible. """
         pass
 
-    # -- _pre_sample and _pre_warmup methods: can be overridden by subclasses --
-    def _pre_sample(self):
-        """ Any code that needs to be run before sampling. """
-        pass
-
-    def _pre_warmup(self):
-        """ Any code that needs to be run before warmup. """
+    @abstractmethod
+    def _initialize(self):
+        """ Subclass specific sampler initialization. Called during the initialization of the sampler which is done before sampling starts. """
         pass
 
     # ------------ Public attributes ------------
     @property
-    def dim(self):
+    def dim(self) -> int:
         """ Dimension of the target density. """
         return self.target.dim
 
     @property
-    def geometry(self):
+    def geometry(self) -> cuqi.geometry.Geometry:
         """ Geometry of the target density. """
         return self.target.geometry
     
@@ -101,39 +136,49 @@ class SamplerNew(ABC):
     def target(self, value):
         """ Set the target density. Runs validation of the target. """
         self._target = value
-        self.validate_target()
-
-    @property
-    def current_point(self):
-        """ The current point of the sampler. """
-        return self._current_point
-    
-    @current_point.setter
-    def current_point(self, value):
-        """ Set the current point of the sampler. """
-        self._current_point = value
+        if self._target is not None:
+            self.validate_target()
 
     # ------------ Public methods ------------
-
     def get_samples(self) -> Samples:
         """ Return the samples. The internal data-structure for the samples is a dynamic list so this creates a copy. """
         return Samples(np.array(self._samples).T, self.target.geometry)
     
-    def reset(self): # TODO. Issue here. Current point is not reset, and initial point is lost with this reset.
-        self._samples.clear()
-        self._acc.clear()
+    def reinitialize(self):
+        """ Re-initialize the sampler. This clears the state and history and initializes the sampler again by setting state and history to their original values. """
+
+        # Loop over state and reset to None
+        for key in self._STATE_KEYS:
+            setattr(self, key, None)
+
+        # Loop over history and reset to None
+        for key in self._HISTORY_KEYS:
+            setattr(self, key, None)
+
+        self._is_initialized = False
+
+        self.initialize()
     
     def save_checkpoint(self, path):
         """ Save the state of the sampler to a file. """
 
+        self._ensure_initialized()
+
         state = self.get_state()
 
+        # Convert all CUQIarrays to numpy arrays since CUQIarrays do not get pickled correctly
+        for key, value in state['state'].items():
+            if isinstance(value, cuqi.array.CUQIarray):
+                state['state'][key] = value.to_numpy()
+
         with open(path, 'wb') as handle:
             pkl.dump(state, handle, protocol=pkl.HIGHEST_PROTOCOL)
 
     def load_checkpoint(self, path):
         """ Load the state of the sampler from a file. """
 
+        self._ensure_initialized()
+
         with open(path, 'rb') as handle:
             state = pkl.load(handle)
 
@@ -155,12 +200,14 @@ class SamplerNew(ABC):
 
         """
 
+        self._ensure_initialized()
+
         # Initialize batch handler
         if batch_size > 0:
             batch_handler = _BatchHandler(batch_size, sample_path)
 
         # Any code that needs to be run before sampling
-        self._pre_sample()
+        if hasattr(self, "_pre_sample"): self._pre_sample()
 
         # Draw samples
         for _ in progressbar( range(Ns) ):
@@ -195,10 +242,12 @@ class SamplerNew(ABC):
 
         """
 
+        self._ensure_initialized()
+
         tune_interval = max(int(tune_freq * Nb), 1)
 
         # Any code that needs to be run before warmup
-        self._pre_warmup()
+        if hasattr(self, "_pre_warmup"): self._pre_warmup()
 
         # Draw warmup samples with tuning
         for idx in progressbar(range(Nb)):
@@ -306,20 +355,60 @@ class SamplerNew(ABC):
                 raise ValueError(f"Key {key} not recognized in history dictionary of sampler {self.__class__.__name__}.")
 
     # ------------ Private methods ------------
-
     def _call_callback(self, sample, sample_index):
         """ Calls the callback function. Assumes input is sample and sample index"""
         if self.callback is not None:
             self.callback(sample, sample_index)
 
+    def _validate_initialization(self):
+        """ Validate the initialization of the sampler by checking all state and history keys are set. """
+
+        for key in self._STATE_KEYS:
+            if getattr(self, key) is None:
+                raise ValueError(f"Sampler state key {key} is not set after initialization.")
+
+        for key in self._HISTORY_KEYS:
+            if getattr(self, key) is None:
+                raise ValueError(f"Sampler history key {key} is not set after initialization.")
+            
+    def _ensure_initialized(self):
+        """ Ensure the sampler is initialized. If not initialize it. """
+        if not self._is_initialized:
+            self.initialize()
+            
+    @property
+    def _default_initial_point(self):
+        """ Return the default initial point for the sampler. Defaults to an array of ones. """
+        return np.ones(self.dim)
+    
+    def __repr__(self):
+        """ Return a string representation of the sampler. """
+        if self.target is None:
+            return f"Sampler: {self.__class__.__name__} \n Target: None"
+        state = self.get_state()
+        msg = f" Sampler: \n\t {self.__class__.__name__} \n Target: \n \t {self.target} \n Current state: \n"
+        # Sort keys alphabetically
+        keys = sorted(state['state'].keys())
+        # Put _ in the end
+        keys = [key for key in keys if key[0] != '_'] + [key for key in keys if key[0] == '_']
+        for key in keys:
+            value = state['state'][key]
+            msg += f"\t {key}: {value} \n"
+        return msg
 
 class ProposalBasedSamplerNew(SamplerNew, ABC):
     """ Abstract base class for samplers that use a proposal distribution. """
 
     _STATE_KEYS = SamplerNew._STATE_KEYS.union({'current_target_logd', 'scale'})
 
-    def __init__(self, target, proposal=None, scale=1, **kwargs):
-        """ Initializer for proposal based samplers. 
+    def __init__(self, target=None, proposal=None, scale=1, **kwargs):
+        """ Initializer for abstract base class for samplers that use a proposal distribution.
+
+        Any subclassing samplers should simply store input parameters as part of the __init__ method.
+
+        Initialization of the sampler should be done in the _initialize method.
+
+        See :class:`SamplerNew` for additional details.
 
         Parameters
         ----------
@@ -338,35 +427,62 @@ class ProposalBasedSamplerNew(SamplerNew, ABC):
         """
 
         super().__init__(target, **kwargs)
+        self.proposal = proposal
+        self.initial_scale = scale
+
+    def initialize(self):
+        """ Initialize the sampler by setting and allocating the state and history before sampling starts. """
+
+        if self._is_initialized:
+            raise ValueError("Sampler is already initialized.")
+        
+        if self.target is None:
+            raise ValueError("Cannot initialize sampler without a target density.")
+        
+        # Default values
+        if self.initial_point is None:
+            self.initial_point = self._default_initial_point
+
+        if self.proposal is None:
+            self.proposal = self._default_proposal
 
+        # State variables
         self.current_point = self.initial_point
+        self.scale = self.initial_scale
+
         self.current_target_logd = self.target.logd(self.current_point)
-        self.proposal = proposal
-        self.scale = scale
 
-        self._acc = [ 1 ] # TODO. Check
+        # History variables
+        self._samples = []
+        self._acc = [ 1 ] # TODO. Check if we need to put 1 here.
 
-    @property 
-    def proposal(self):
-        return self._proposal 
+        self._initialize() # Subclass specific initialization
+
+        self._validate_initialization()
+
+        self._is_initialized = True
+
+    @abstractmethod
+    def validate_proposal(self):
+        """ Validate the proposal distribution. """
+        pass     
 
-    @proposal.setter 
-    def proposal(self, value):
-        self._proposal = value
+    @property
+    def _default_proposal(self):
+        """ Return the default proposal distribution. Defaults to a Gaussian distribution with zero mean and unit variance. """
+        return cuqi.distribution.Gaussian(np.zeros(self.dim), 1)
 
     @property
-    def geometry(self): # TODO. Check if we can refactor this
-        geom1, geom2 = None, None
-        if hasattr(self, 'proposal') and hasattr(self.proposal, 'geometry') and self.proposal.geometry.par_dim is not None:
-            geom1=  self.proposal.geometry
-        if hasattr(self, 'target') and hasattr(self.target, 'geometry') and self.target.geometry.par_dim is not None:
-            geom2 = self.target.geometry
-        if not isinstance(geom1,cuqi.geometry._DefaultGeometry) and geom1 is not None:
-            return geom1
-        elif not isinstance(geom2,cuqi.geometry._DefaultGeometry) and geom2 is not None: 
-            return geom2
-        else:
-            return cuqi.geometry._DefaultGeometry(self.dim)
+    def proposal(self):
+        """ The proposal distribution. """
+        return self._proposal
+    
+    @proposal.setter
+    def proposal(self, proposal):
+        """ Set the proposal distribution. """
+        self._proposal = proposal
+        if self._proposal is not None:
+            self.validate_proposal()
 
 
 class _BatchHandler: