<commit_before>from numpy import dot
from numpy.random import normal
from numpy.linalg import solve
from scipy.linalg import cholesky, cho_solve
import numpy as np

def quad_potential(C, is_cov):
    partial_check_positive_definite(C)
    if C.ndim == 1:
        if is_cov:
            return ElemWiseQuadPotential(C)
        else: 
            return ElemWiseQuadPotential(1./C)
    else:
        if is_cov:
            return QuadPotential(C)
        else :
            return QuadPotential_Inv(C) 

def partial_check_positive_definite(C):
    """Simple but partial check for Positive Definiteness"""
    if C.ndim == 1:
        d = C 
    else :
        d = np.diag(C)
    i, = np.nonzero(np.logical_or(np.isnan(d), d<=0))

    if len(i):
        raise PositiveDefiniteError("Simple check failed. Diagonal contains negatives", i)

class PositiveDefiniteError(ValueError):
    def __init__(self, msg, idx): 
        self.idx = idx
        self.msg = msg
    def __str__(self):
        return "Scaling is not positive definite. " + self.msg + ". Check indexes " + str(self.idx)

class ElemWiseQuadPotential(object):
    def __init__(self, v):
        s = v **.5

        self.s = s
        self.inv_s = 1./s
        self.v = v 
    def velocity(self, x):
        return self.v* x
    def random(self):
        return normal(size = self.s.shape)* self.inv_s
    def energy(self, x):
        return .5*dot(x, self.v*x)


class QuadPotential_Inv(object):
    def __init__(self, A):
        self.L = cholesky(A, lower = True)
        
    def velocity(self, x ):
        return cho_solve((self.L, True), x)
        
    def random(self):
        n = normal(size = self.L.shape[0])
        return dot(self.L, n)
    
    def energy(self, x):
        L1x = solve(self.L, x)
        return .5 * dot(L1x.T, L1x)


class QuadPotential(object):
    def __init__(self, A):
        self.A = A
        self.L = cholesky(A, lower = True)
    
    def velocity(self, x):
        return dot(self.A, x)
    
    def random(self):
        n = normal(size = self.L.shape[0])
        return solve(self.L.T, n)<commit_msg>allow quad_potential and thus hmc to take sparse matrix and quad potentials<commit_after>from numpy import dot
from numpy.random import normal
from numpy.linalg import solve
from scipy.linalg import cholesky, cho_solve
from scipy.sparse import issparse

import numpy as np

def quad_potential(C, is_cov):
    if isquadpotential(C):
        return C

    if issparse(C):
        return QuadPotential_SparseInv(C)
    
    partial_check_positive_definite(C)
    if C.ndim == 1:
        if is_cov:
            return ElemWiseQuadPotential(C)
        else: 
            return ElemWiseQuadPotential(1./C)
    else:
        if is_cov:
            return QuadPotential(C)
        else :
            return QuadPotential_Inv(C) 

def partial_check_positive_definite(C):
    """Simple but partial check for Positive Definiteness"""
    if C.ndim == 1:
        d = C 
    else :
        d = np.diag(C)
    i, = np.nonzero(np.logical_or(np.isnan(d), d<=0))

    if len(i):
        raise PositiveDefiniteError("Simple check failed. Diagonal contains negatives", i)

class PositiveDefiniteError(ValueError):
    def __init__(self, msg, idx): 
        self.idx = idx
        self.msg = msg
    def __str__(self):
        return "Scaling is not positive definite. " + self.msg + ". Check indexes " + str(self.idx)

def isquadpotential(o): 
    return all(hasattr(o, attr) for attr in ["velocity", "random", "energy"])

class ElemWiseQuadPotential(object):
    def __init__(self, v):
        s = v **.5

        self.s = s
        self.inv_s = 1./s
        self.v = v 
    def velocity(self, x):
        return self.v* x
    def random(self):
        return normal(size = self.s.shape)* self.inv_s
    def energy(self, x):
        return .5*dot(x, self.v*x)

class QuadPotential_Inv(object):
    def __init__(self, A):
        self.L = cholesky(A, lower = True)
        
    def velocity(self, x ):
        return cho_solve((self.L, True), x)
        
    def random(self):
        n = normal(size = self.L.shape[0])
        return dot(self.L, n)
    
    def energy(self, x):
        L1x = solve(self.L, x)
        return .5 * dot(L1x.T, L1x)


class QuadPotential(object):
    def __init__(self, A):
        self.A = A
        self.L = cholesky(A, lower = True)
    
    def velocity(self, x):
        return dot(self.A, x)
    
    def random(self):
        n = normal(size = self.L.shape[0])
        return solve(self.L.T, n)
{'fraction_non_alphanumeric': 0.11271036315323295, 'fraction_numerical': 0.005757307351638618, 'mean_word_length': 2.217236467236467, 'pattern_counts': {'":': 0, '<': 5, '<?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': '7467547', 'n_tokens_mistral': 1609, 'n_tokens_neox': 1528, 'n_words': 405}