<commit_before>'''Kalman filtering / smoothing module.'''


import abc
import collections

import numpy as np
import scipy.linalg


class DTKalmanFilterBase(metaclass=abc.ABCMeta):
    '''Discrete-time Kalman filter/smoother abstract base class.'''
    
    def __init__(self, model, mean, cov, **options):
        self.model = model
        '''The underlying system model.'''
        
        self.mean = np.asarray(mean, dtype=float)
        '''The latest filtered state mean.'''
        assert (self.model.nx,) == self.mean.shape
        
        self.cov = np.asarray(cov, dtype=float)
        '''The latest filtered state covariance.'''
        assert (self.model.nx, self.model.nx) == self.cov.shape
        
        self.options = options
        '''Filter options.'''
        
        self.k = 0
        '''The latest working sample number.'''
        
        if self.save_history:
            self.pred_mean = [self.mean]
            '''Predicted state mean history.'''
            
            self.pred_cov = [self.cov]
            '''Predicted state covariance history.'''
            
            self.corr_mean = [self.mean]
            '''Corrected state mean history.'''
            
            self.corr_cov = [self.cov]
            '''Corrected state covariance history.'''
    
        if self.save_pred_xcov:
            self.pred_xcov = []
            '''State prediction cross-covariance history.'''
        
        if self.save_likelihood:
            self.loglikelihood = 0
            '''Measurement log-likelihood.'''
 
    @property
    def save_history(self):
        '''Whether to save the filtering history or just the latest values.'''
        return self.options.get('save_history', True)

    @property
    def save_likelihood(self):
        '''Whether to save the measurement likelihood.'''
        return self.options.get('save_likelihood', True)

    @property
    def save_pred_xcov(self):
        '''Whether to save the prediction cross-covariance matrices.'''
        return self.save_history and self.options.get('save_pred_xcov', True)
    
    def _save_prediction(self, mean, cov):
        '''Save the prection data and update time index.'''
        self.mean = mean
        self.cov = cov
        self.k += 1
        
        if self.save_history:
            self.pred_mean.append(mean)
            self.pred_cov.append(cov)
            self.corr_mean.append(mean)
            self.corr_cov.append(cov)
    
    def _save_correction(self, mean, cov):
        '''Save the correction data.'''
        self.mean = mean
        self.cov = cov
        
        if self.save_history:
            self.corr_mean[-1] = mean
            self.corr_cov[-1] = cov
    
    @abc.abstractmethod
    def predict(self, u=[]):
        '''Predict the state distribution at the next time index.'''
        raise NotImplementedError("Pure abstract method.")
    
    @abc.abstractmethod
    def correct(self, y, u=[]):
        '''Correct the state distribution, given the measurement vector.'''
        raise NotImplementedError("Pure abstract method.")
    
    def filter(self, y, u=None):
        for k, yk in enumerate(y[:-1]):
            uk = [] if u is None else u[k]
            self.correct(yk, uk)
            self.predict(uk)
        
        self.correct(y[-1], [] if u is None else u[-1])


class UnscentedTransform:
    
    def __init__(self, **options):
        self.options = options
    
    @property
    def kappa(self):
        return self.options.get('kappa', 0.0)
    
    @property
    def sqrt(self):
        return self.options.get('sqrt', 'svd')
    
    def _gen_sigma_points(self, mean, cov):
        '''Generate sigma-points deviations and weights.'''
        mean = np.asarray(mean)
        cov = np.asarray(cov)
        n = len(mean)
        kappa = self.kappa
        sqrt = self.sqrt
        assert n + kappa != 0
        
        if sqrt == 'svd':
            [U, s, Vh] = scipy.linalg.svd((n + kappa) * cov)
            cov_sqrt = (U * np.sqrt(s)).T
        elif sqrt == 'cholesky':
            cov_sqrt = scipy.linalg.cholesky((n + kappa) * cov, lower=False)
        elif isinstance(sqrt, collections.Callable):
            cov_sqrt = sqrt((n + kappa) * cov)
        else:
            raise ValueError("Unknown option for 'sqrt' argument.")
        
        dev = np.vstack([cov_sqrt, -cov_sqrt])
        weights = np.repeat(0.5 / (n + kappa), 2 * n)
        
        if kappa != 0:
            dev = np.vstack([np.zeros_like(mean), dev])
            weights = np.hstack([kappa / (n + kappa), weights])
        
        self.input_dev = dev
        self.weights = weights
        self.input_sigma = dev + mean
        
        return (self.input_sigma, weights)
    
    def unscented_transform(self, f, mean, cov):
        [input_sigma, weights] = self._gen_sigma_points(mean, cov)
        
        output_sigma = [f(x) for x in input_sigma]
        output_mean = np.dot(weights, output_sigma)
        output_dev = output_sigma - output_mean
        output_cov = np.einsum('ki,kj,k', output_dev, output_dev, weights)
        
        self.output_dev = output_dev
        return (output_mean, output_cov)
    
    def transform_xcov(self):
        try:
            input_dev = self.input_dev
            output_dev = self.output_dev
            weights = self.weights
        except AttributeError:
            raise RuntimeError(
                "Unscented transform must be done before requesting xcov."
            )
        
        return np.einsum('ki,kj,k', input_dev, output_dev, weights)

# Transpose dev and sigma
# Transpose sqrt
# Apply f(sigma) vectorized


class DTUnscentedPredictor(DTKalmanFilterBase, UnscentedTransform):
    <commit_msg>kalman: transposed the sqrt of the unscented transform.<commit_after>'''Kalman filtering / smoothing module.'''


import abc
import collections

import numpy as np
import scipy.linalg


class DTKalmanFilterBase(metaclass=abc.ABCMeta):
    '''Discrete-time Kalman filter/smoother abstract base class.'''
    
    def __init__(self, model, mean, cov, **options):
        self.model = model
        '''The underlying system model.'''
        
        self.mean = np.asarray(mean, dtype=float)
        '''The latest filtered state mean.'''
        assert (self.model.nx,) == self.mean.shape
        
        self.cov = np.asarray(cov, dtype=float)
        '''The latest filtered state covariance.'''
        assert (self.model.nx, self.model.nx) == self.cov.shape
        
        self.options = options
        '''Filter options.'''
        
        self.k = 0
        '''The latest working sample number.'''
        
        if self.save_history:
            self.pred_mean = [self.mean]
            '''Predicted state mean history.'''
            
            self.pred_cov = [self.cov]
            '''Predicted state covariance history.'''
            
            self.corr_mean = [self.mean]
            '''Corrected state mean history.'''
            
            self.corr_cov = [self.cov]
            '''Corrected state covariance history.'''
    
        if self.save_pred_xcov:
            self.pred_xcov = []
            '''State prediction cross-covariance history.'''
        
        if self.save_likelihood:
            self.loglikelihood = 0
            '''Measurement log-likelihood.'''
 
    @property
    def save_history(self):
        '''Whether to save the filtering history or just the latest values.'''
        return self.options.get('save_history', True)

    @property
    def save_likelihood(self):
        '''Whether to save the measurement likelihood.'''
        return self.options.get('save_likelihood', True)

    @property
    def save_pred_xcov(self):
        '''Whether to save the prediction cross-covariance matrices.'''
        return self.save_history and self.options.get('save_pred_xcov', True)
    
    def _save_prediction(self, mean, cov):
        '''Save the prection data and update time index.'''
        self.mean = mean
        self.cov = cov
        self.k += 1
        
        if self.save_history:
            self.pred_mean.append(mean)
            self.pred_cov.append(cov)
            self.corr_mean.append(mean)
            self.corr_cov.append(cov)
    
    def _save_correction(self, mean, cov):
        '''Save the correction data.'''
        self.mean = mean
        self.cov = cov
        
        if self.save_history:
            self.corr_mean[-1] = mean
            self.corr_cov[-1] = cov
    
    @abc.abstractmethod
    def predict(self, u=[]):
        '''Predict the state distribution at the next time index.'''
        raise NotImplementedError("Pure abstract method.")
    
    @abc.abstractmethod
    def correct(self, y, u=[]):
        '''Correct the state distribution, given the measurement vector.'''
        raise NotImplementedError("Pure abstract method.")
    
    def filter(self, y, u=None):
        for k, yk in enumerate(y[:-1]):
            uk = [] if u is None else u[k]
            self.correct(yk, uk)
            self.predict(uk)
        
        self.correct(y[-1], [] if u is None else u[-1])


class UnscentedTransform:
    
    def __init__(self, **options):
        self.options = options
    
    @property
    def kappa(self):
        return self.options.get('kappa', 0.0)
    
    @property
    def sqrt(self):
        return self.options.get('sqrt', 'svd')
    
    def _gen_sigma_points(self, mean, cov):
        '''Generate sigma-points deviations and weights.'''
        mean = np.asarray(mean)
        cov = np.asarray(cov)
        n = len(mean)
        kappa = self.kappa
        sqrt = self.sqrt
        assert n + kappa != 0
        
        if sqrt == 'svd':
            [U, s, Vh] = scipy.linalg.svd((n + kappa) * cov)
            cov_sqrt = U * np.sqrt(s)
        elif sqrt == 'cholesky':
            cov_sqrt = scipy.linalg.cholesky((n + kappa) * cov, lower=True)
        elif isinstance(sqrt, collections.Callable):
            cov_sqrt = sqrt((n + kappa) * cov)
        else:
            raise ValueError("Unknown option for 'sqrt' argument.")
        
        dev = np.hstack((cov_sqrt, -cov_sqrt))
        weights = np.repeat(0.5 / (n + kappa), 2 * n)
        
        if kappa != 0:
            dev = np.hstack((np.zeros_like(mean), dev))
            weights = np.hstack([kappa / (n + kappa), weights])
        
        self.input_dev = dev
        self.weights = weights
        self.input_sigma = dev + mean[:, None]
        
        return (self.input_sigma, weights)
    
    def unscented_transform(self, f, mean, cov):
        [input_sigma, weights] = self._gen_sigma_points(mean, cov)
        
        output_sigma = f(input_sigma)
        output_mean = np.dot(output_sigma, weights)
        output_dev = output_sigma - output_mean[:, None]
        output_cov = np.einsum('ik,jk,k', output_dev, output_dev, weights)
        
        self.output_dev = output_dev
        return (output_mean, output_cov)
    
    def transform_xcov(self):
        try:
            input_dev = self.input_dev
            output_dev = self.output_dev
            weights = self.weights
        except AttributeError:
            raise RuntimeError(
                "Unscented transform must be done before requesting xcov."
            )
        
        return np.einsum('ik,jk,k', input_dev, output_dev, weights)


class DTUnscentedPredictor(DTKalmanFilterBase, UnscentedTransform):
    
{'fraction_non_alphanumeric': 0.11340929009640666, 'fraction_numerical': 0.0026292725679228747, 'mean_word_length': 1.9251473981030505, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '306457', 'n_tokens_mistral': 3430, 'n_tokens_neox': 3316, 'n_words': 852}