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

cuqi/experimental/mcmc/_sampler.py

@@ -1,582 +0,0 @@
-from abc import ABC, abstractmethod
-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.
-
-    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. """
-
-    _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)`.
-        """
-
-        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
-
-    @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) -> int:
-        """ Dimension of the target density. """
-        return self.target.dim
-
-    @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):
-        """ Set the target density. Runs validation of the target. """
-        self._target = value
-        if self._target is not None:
-            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
-
-    # ------------ 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. """
-
-        # 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)
-
-        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)
-                
-        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:
-                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:
-            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
-
-class ProposalBasedSampler(Sampler, ABC):
-    """ Abstract base class for samplers that use a proposal distribution. """
-
-    _STATE_KEYS = Sampler._STATE_KEYS.union({'current_target_logd', 'scale'})
-
-    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:`Sampler` for additional details.
-
-        Parameters
-        ----------
-        target : cuqi.density.Density
-            The target density.
-
-        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):
-        """ 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:
-    """ 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 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()