CUQIpy 1.4.0.post0.dev13__py3-none-any.whl → 1.4.0.post0.dev45__py3-none-any.whl

cuqi/sampler/_sampler.py

@@ -1,177 +1,582 @@
 from abc import ABC, abstractmethod
-import sys
+import os
 import numpy as np
+import pickle as pkl
+import warnings
 import cuqi
 from cuqi.samples import Samples
 
+try:
+    from tqdm import tqdm
+except ImportError:
+    def tqdm(iterable, **kwargs):
+        warnings.warn("Module mcmc: tqdm not found. Install tqdm to get sampling progress.")
+        return iterable
+
 class Sampler(ABC):
+    """ Abstract base class for all samplers.
+    
+    Provides a common interface for all samplers. The interface includes methods for sampling, warmup and getting the samples in an object oriented way.
 
-    def __init__(self, target, x0=None, dim=None, callback=None):
+    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.
 
-        self._dim = dim
-        if hasattr(target,'dim'): 
-            if self._dim is None:
-                self._dim = target.dim 
-            elif self._dim != target.dim:
-                raise ValueError("'dim' need to be None or equal to 'target.dim'") 
-        elif x0 is not None:
-            self._dim = len(x0)
+    The sampler maintains sets of state and history keys, which are used for features like checkpointing and resuming sampling.
 
-        self.target = target
+    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.
 
-        if x0 is None:
-            x0 = np.ones(self.dim)
-        self.x0 = x0
+    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.
 
-        self.callback = callback
+    Subclasses should ensure that any new variables that are updated in a Markov Monte Carlo step are added to the state or history keys.
 
-    def step(self, x):
-        """
-        Perform a single MCMC step
-        """
-        # Currently a hack to get step method for any sampler
-        self.x0 = x
-        return self.sample(2).samples[:,-1]
+    Saving and loading checkpoints saves and loads the state of the sampler (not the history).
 
-    def step_tune(self, x, *args, **kwargs):
-        """
-        Perform a single MCMC step and tune the sampler. This is used during burn-in.
-        """
-        # Currently a hack to get step method for any sampler
-        out = self.step(x)
-        self.tune(*args, *kwargs)
-        return out
+    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.
+
+    """
 
-    def tune(self):
+    _STATE_KEYS = {'current_point'}
+    """ Set of keys for the state dictionary. """
+
+    _HISTORY_KEYS = {'_samples', '_acc'}
+    """ Set of keys for the history dictionary. """
+
+    def __init__(self, target:cuqi.density.Density=None, 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
+        ----------
+        target : cuqi.density.Density
+            The target density.
+
+        initial_point : array-like, optional
+            The initial point for the sampler. If not given, the sampler will choose an initial point.
+
+        callback : callable, optional
+            A function that will be called after each sampling step. It can be useful for monitoring the sampler during sampling.
+            The function should take three arguments: the sampler object, the index of the current sampling step, the total number of requested samples. The last two arguments are integers. An example of the callback function signature is: `callback(sampler, sample_index, num_of_samples)`.
         """
-        Tune the sampler parameters.
+
+        self.target = target
+        self.initial_point = initial_point
+        self.callback = callback
+        self._is_initialized = False
+
+    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._get_default_initial_point(self.dim)
+
+        # 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. """
+        pass
+
+    @abstractmethod
+    def tune(self, skip_len, update_count):
+        """ Tune the parameters of the sampler. This method is called after each step of the warmup phase.
+
+        Parameters
+        ----------
+        skip_len : int
+            Defines the number of steps in between tuning (i.e. the tuning interval).
+
+        update_count : int
+            The number of times tuning has been performed. Can be used for internal bookkeeping.
+        
         """
         pass
 
+    @abstractmethod
+    def validate_target(self):
+        """ 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
 
-    @property
-    def geometry(self):
-        if hasattr(self, 'target') and hasattr(self.target, 'geometry'):
-            geom =  self.target.geometry
-        else:
-            geom = cuqi.geometry._DefaultGeometry1D(self.dim)
-        return geom
+    @abstractmethod
+    def _initialize(self):
+        """ Subclass specific sampler initialization. Called during the initialization of the sampler which is done before sampling starts. """
+        pass
 
-    @property 
-    def target(self):
-        return self._target 
+    # ------------ Public attributes ------------
+    @property
+    def dim(self) -> int:
+        """ Dimension of the target density. """
+        return self.target.dim
 
-    @target.setter 
+    @property
+    def geometry(self) -> cuqi.geometry.Geometry:
+        """ Geometry of the target density. """
+        return self.target.geometry
+    
+    @property
+    def target(self) -> cuqi.density.Density:
+        """ Return the target density. """
+        return self._target
+    
+    @target.setter
     def target(self, value):
-        if  not isinstance(value, cuqi.distribution.Distribution) and callable(value):
-            # obtain self.dim
-            if self.dim is not None:
-                dim = self.dim
-            else:
-                raise ValueError(f"If 'target' is a lambda function, the parameter 'dim' need to be specified when initializing {self.__class__}.")
+        """ Set the target density. Runs validation of the target. """
+        self._target = value
+        if self._target is not None:
+            self.validate_target()
 
-            # set target
-            self._target = cuqi.distribution.UserDefinedDistribution(logpdf_func=value, dim = dim)
+    @property
+    def current_point(self):
+        """ The current point of the sampler. """
+        return self._current_point
 
-        elif isinstance(value, cuqi.distribution.Distribution):
-            self._target = value
-        else:
-            raise ValueError("'target' need to be either a lambda function or of type 'cuqi.distribution.Distribution'")
+    @current_point.setter
+    def current_point(self, value):
+        """ Set the current point of the sampler. """
+        self._current_point = value
 
+    # ------------ 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 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. """
 
-    @property
-    def dim(self):
-        if hasattr(self,'target') and hasattr(self.target,'dim'):
-            self._dim = self.target.dim 
-        return self._dim
+        # 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. """
 
-    def sample(self,N,Nb=0):
-        # Get samples from the samplers sample method
-        result = self._sample(N,Nb)
-        return self._create_Sample_object(result,N+Nb)
-
-    def sample_adapt(self,N,Nb=0):
-        # Get samples from the samplers sample method
-        result = self._sample_adapt(N,Nb)
-        return self._create_Sample_object(result,N+Nb)
-
-    def _create_Sample_object(self,result,N):
-        loglike_eval = None
-        acc_rate = None
-        if isinstance(result,tuple):
-            #Unpack samples+loglike+acc_rate
-            s = result[0]
-            if len(result)>1: loglike_eval = result[1]
-            if len(result)>2: acc_rate = result[2]
-            if len(result)>3: raise TypeError("Expected tuple of at most 3 elements from sampling method.")
-        else:
-            s = result
+        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)
+
+        self.set_state(state)
+
+    def sample(self, Ns, Nt=1, batch_size=0, sample_path='./CUQI_samples/') -> 'Sampler':
+        """ Sample Ns samples from the target density.
+
+        Parameters
+        ----------
+        Ns : int
+            The number of samples to draw.
+        
+        Nt : int, optional, default=1
+            The thinning interval. If Nt >= 1, every Nt'th sample is stored. The larger Nt, the fewer samples are stored.
+
+        batch_size : int, optional
+            The batch size for saving samples to disk. If 0, no batching is used. If positive, samples are saved to disk in batches of the specified size.
+
+        sample_path : str, optional
+            The path to save the samples. If not specified, the samples are saved to the current working directory under a folder called 'CUQI_samples'.
+
+        """
+        self._ensure_initialized()
+
+        # Initialize batch handler
+        if batch_size > 0:
+            batch_handler = _BatchHandler(batch_size, sample_path)
+
+        # Draw samples
+        pbar = tqdm(range(Ns), "Sample: ")
+        for idx in pbar:
+            
+            # Perform one step of the sampler
+            acc = self.step()
+
+            # Store samples
+            self._acc.append(acc)
+            if (Nt > 0) and ((idx + 1) % Nt == 0):
+                self._samples.append(self.current_point)
+
+            # display acc rate at progress bar
+            pbar.set_postfix_str(f"acc rate: {np.mean(self._acc[-1-idx:]):.2%}")
+
+            # Add sample to batch
+            if batch_size > 0:
+                batch_handler.add_sample(self.current_point)
+
+            # Call callback function if specified            
+            self._call_callback(idx, Ns)
                 
-        #Store samples in cuqi samples object if more than 1 sample
-        if N==1:
-            if len(s) == 1 and isinstance(s,np.ndarray): #Extract single value from numpy array
-                s = s.ravel()[0]
+        return self
+    
+
+    def warmup(self, Nb, Nt=1, tune_freq=0.1) -> 'Sampler':
+        """ Warmup the sampler by drawing Nb samples.
+
+        Parameters
+        ----------
+        Nb : int
+            The number of samples to draw during warmup.
+
+        Nt : int, optional, default=1
+            The thinning interval. If Nt >= 1, every Nt'th sample is stored. The larger Nt, the fewer samples are stored.
+            
+        tune_freq : float, optional
+            The frequency of tuning. Tuning is performed every tune_freq*Nb samples.
+
+        """
+
+        self._ensure_initialized()
+
+        tune_interval = max(int(tune_freq * Nb), 1)
+
+        # Draw warmup samples with tuning
+        pbar = tqdm(range(Nb), "Warmup: ")
+        for idx in pbar:
+
+            # Perform one step of the sampler
+            acc = self.step()
+
+            # Tune the sampler at tuning intervals
+            if (idx + 1) % tune_interval == 0:
+                self.tune(tune_interval, idx // tune_interval) 
+
+            # Store samples
+            self._acc.append(acc)
+            if (Nt > 0) and ((idx + 1) % Nt == 0):
+                self._samples.append(self.current_point)
+
+            # display acc rate at progress bar
+            pbar.set_postfix_str(f"acc rate: {np.mean(self._acc[-1-idx:]):.2%}")
+
+            # Call callback function if specified
+            self._call_callback(idx, Nb)
+
+        return self
+    
+    def get_state(self) -> dict:
+        """ Return the state of the sampler. 
+
+        The state is used when checkpointing the sampler.
+
+        The state of the sampler is a dictionary with keys 'metadata' and 'state'.
+        The 'metadata' key contains information about the sampler type.
+        The 'state' key contains the state of the sampler.
+
+        For example, the state of a "MH" sampler could be:
+
+        state = {
+            'metadata': {
+                'sampler_type': 'MH'
+            },
+            'state': {
+                'current_point': np.array([...]),
+                'current_target_logd': -123.45,
+                'scale': 1.0,
+                ...
+            }
+        }
+        """   
+        state = {
+            'metadata': {
+                'sampler_type': self.__class__.__name__
+            },
+            'state': {
+                key: getattr(self, key) for key in self._STATE_KEYS
+            }
+        }
+        return state
+
+    def set_state(self, state: dict):
+        """ Set the state of the sampler.
+
+        The state is used when loading the sampler from a checkpoint.
+
+        The state of the sampler is a dictionary with keys 'metadata' and 'state'.
+
+        For example, the state of a "MH" sampler could be:
+
+        state = {
+            'metadata': {
+                'sampler_type': 'MH'
+            },
+            'state': {
+                'current_point': np.array([...]),
+                'current_target_logd': -123.45,
+                'scale': 1.0,
+                ...
+            }
+        }
+        """
+        if state['metadata']['sampler_type'] != self.__class__.__name__:
+            raise ValueError(f"Sampler type in state dictionary ({state['metadata']['sampler_type']}) does not match the type of the sampler ({self.__class__.__name__}).")
+        
+        for key, value in state['state'].items():
+            if key in self._STATE_KEYS:
+                setattr(self, key, value)
             else:
-                s = s.flatten()
+                raise ValueError(f"Key {key} not recognized in state dictionary of sampler {self.__class__.__name__}.")
+            
+    def get_history(self) -> dict:
+        """ Return the history of the sampler. """
+        history = {
+            'metadata': {
+                'sampler_type': self.__class__.__name__
+            },
+            'history': {
+                key: getattr(self, key) for key in self._HISTORY_KEYS
+            }
+        }
+        return history
+    
+    def set_history(self, history: dict):
+        """ Set the history of the sampler. """
+        if history['metadata']['sampler_type'] != self.__class__.__name__:
+            raise ValueError(f"Sampler type in history dictionary ({history['metadata']['sampler_type']}) does not match the type of the sampler ({self.__class__.__name__}).")
+        
+        for key, value in history['history'].items():
+            if key in self._HISTORY_KEYS:
+                setattr(self, key, value)
+            else:
+                raise ValueError(f"Key {key} not recognized in history dictionary of sampler {self.__class__.__name__}.")
+
+    # ------------ Private methods ------------
+    def _call_callback(self, sample_index, num_of_samples):
+        """ Calls the callback function. Assumes input is sampler, sample index, and total number of samples """
+        if self.callback is not None:
+            self.callback(self, sample_index, num_of_samples)
+
+    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()
+            
+    def _get_default_initial_point(self, dim):
+        """ Return the default initial point for the sampler. Defaults to an array of ones. """
+        return np.ones(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"
         else:
-            s = Samples(s, self.geometry)#, geometry = self.geometry)
-            s.loglike_eval = loglike_eval
-            s.acc_rate = acc_rate
-        return s
+            msg = f"Sampler: {self.__class__.__name__} \n Target: \n \t {self.target} "
+            
+        if self._is_initialized:
+            state = self.get_state()
+            msg += f"\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
 
-    @abstractmethod
-    def _sample(self,N,Nb):
-        pass
+class ProposalBasedSampler(Sampler, ABC):
+    """ Abstract base class for samplers that use a proposal distribution. """
 
-    @abstractmethod
-    def _sample_adapt(self,N,Nb):
-        pass
+    _STATE_KEYS = Sampler._STATE_KEYS.union({'current_target_logd', 'scale'})
 
-    def _print_progress(self,s,Ns):
-        """Prints sampling progress"""
-        if Ns > 2:
-            if (s % (max(Ns//100,1))) == 0:
-                msg = f'Sample {s} / {Ns}'
-                sys.stdout.write('\r'+msg)
-            if s==Ns:
-                msg = f'Sample {s} / {Ns}'
-                sys.stdout.write('\r'+msg+'\n')
-
-    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 __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.
 
-class ProposalBasedSampler(Sampler,ABC):
-    def __init__(self, target,  proposal=None, scale=1, x0=None, dim=None, **kwargs):
-        #TODO: after fixing None dim
-        #if dim is None and hasattr(proposal,'dim'):
-        #    dim = proposal.dim
-        super().__init__(target, x0=x0, dim=dim, **kwargs)
+        Initialization of the sampler should be done in the _initialize method.
 
-        self.proposal =proposal
-        self.scale = scale
+        See :class:`Sampler` for additional details.
 
+        Parameters
+        ----------
+        target : cuqi.density.Density
+            The target density.
 
-    @property 
+        proposal : cuqi.distribution.Distribution, optional
+            The proposal distribution. If not specified, the default proposal is used.
+
+        scale : float, optional
+            The scale parameter for the proposal distribution.
+
+        **kwargs : dict
+            Additional keyword arguments passed to the :class:`Sampler` initializer.
+
+        """
+
+        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._get_default_initial_point(self.dim)
+
+        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)
+
+        # 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
+
+    @abstractmethod
+    def validate_proposal(self):
+        """ Validate the proposal distribution. """
+        pass     
+
+    @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 proposal(self):
-        return self._proposal 
+        """ 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()
+
 
-    @proposal.setter 
-    def proposal(self, value):
-        self._proposal = value
+class _BatchHandler:
+    """ Utility class to handle batching of samples. 
+    
+    If a batch size is specified, this class will save samples to disk in batches of the specified size.
+     
+    This is useful for very large sample sets that do not fit in memory.
+     
+    """
+    
+    def __init__(self, batch_size=0, sample_path='./CUQI_samples/'):
+
+        if batch_size < 0:
+            raise ValueError("Batch size should be a non-negative integer")
+
+        self.sample_path = sample_path
+        self._batch_size = batch_size
+        self.current_batch = []
+        self.num_batches_dumped = 0
 
     @property
-    def geometry(self):
-        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._DefaultGeometry1D(self.dim)
+    def sample_path(self):
+        """ The path to save the samples. """
+        return self._sample_path
+    
+    @sample_path.setter
+    def sample_path(self, value):
+        if not isinstance(value, str):
+            raise TypeError("Sample path must be a string.")
+        normalized_path = value.rstrip('/') + '/'
+        if not os.path.isdir(normalized_path):
+            try:
+                os.makedirs(normalized_path, exist_ok=True)
+            except Exception as e:
+                raise ValueError(f"Could not create directory at {normalized_path}: {e}")
+        self._sample_path = normalized_path
+
+    def add_sample(self, sample):
+        """ Add a sample to the batch if batching. If the batch is full, flush the batch to disk. """
+
+        if self._batch_size <= 0:
+            return  # Batching not used
+
+        self.current_batch.append(sample)
+
+        if len(self.current_batch) >= self._batch_size:
+            self.flush()
+
+    def flush(self):
+        """ Flush the current batch of samples to disk. """
+
+        if not self.current_batch:
+            return  # No samples to flush
+
+        # Save the current batch of samples
+        batch_samples = np.array(self.current_batch)
+        file_path = f'{self.sample_path}batch_{self.num_batches_dumped:04d}.npz'
+        np.savez(file_path, samples=batch_samples, batch_id=self.num_batches_dumped)
+
+        self.num_batches_dumped += 1
+        self.current_batch = []  # Clear the batch after saving
+
+    def finalize(self):
+        """ Finalize the batch handler. Flush any remaining samples to disk. """
+        self.flush()