<commit_before>import os
import sys
import numpy as np
import scipy.ndimage as nd
import tempfile
import pickle

from cbamf import const
from cbamf.mc import samplers, engines, observers

# Linear fit function, because I can
def linear_fit(x, y, sigma=1, N=100, burn=1000):
    from cbamf.states import LinearFit
    poly = np.polyfit(x,y,1)

    s = LinearFit(x, y, sigma=sigma)
    s.state[:2] = poly

    bl = s.explode(s.block_all())
    h = sample_state(s, s.explode(s.block_all()), N=burn, doprint=True, procedure='uniform')
    h = sample_state(s, s.explode(s.block_all()), N=burn, doprint=True, procedure='uniform')

    return s, h.get_histogram()

#=============================================================================
# Sampling methods that run through blocks and sample
#=============================================================================
def sample_state(state, blocks, stepout=1, slicing=True, N=1, doprint=False, procedure='uniform'):
    eng = engines.SequentialBlockEngine(state)
    opsay = observers.Printer()
    ohist = observers.HistogramObserver(block=np.array(blocks).any(axis=0))
    eng.add_samplers([samplers.SliceSampler1D(stepout, block=b, procedure=procedure) for b in blocks])

    eng.add_likelihood_observers(opsay) if doprint else None
    eng.add_state_observers(ohist)

    eng.dosteps(N)
    return ohist

def scan_ll(state, element, size=0.1, N=1000):
    start = state.state[element]

    ll = []
    vals = np.linspace(start-size, start+size, N)
    for val in vals:
        state.update(element, val)
        l = state.loglikelihood()
        ll.append(l)

    state.update(element, start)
    return vals, np.array(ll)

def scan_noise(image, state, element, size=0.01, N=1000):
    start = state.state[element]

    xs, ys = [], []
    for i in xrange(N):
        print i
        test = image + np.random.normal(0, state.sigma, image.shape)
        x,y = sample_ll(test, state, element, size=size, N=300)
        state.update(element, start)
        xs.append(x)
        ys.append(y)

    return xs, ys

def sample_particles(state, stepout=1, start=0):
    print '{:-^39}'.format(' POS / RAD ')
    for particle in xrange(start, state.obj.N):
        if not state.isactive(particle):
            continue

        print particle
        sys.stdout.flush()

        blocks = state.blocks_particle(particle)
        sample_state(state, blocks, stepout=stepout)

    return state.state.copy()

def sample_particle_pos(state, stepout=1, start=0, quiet=False):
    if not quiet:
        print '{:-^39}'.format(' POS ')

    for particle in xrange(start, state.obj.N):
        if not state.isactive(particle):
            continue

        if not quiet:
            print particle
        sys.stdout.flush()

        blocks = state.blocks_particle(particle)[:-1]
        sample_state(state, blocks, stepout=stepout)

    return state.state.copy()

def sample_particle_rad(state, stepout=1, start=0, quiet=False):
    if not quiet:
        print '{:-^39}'.format(' RAD ')

    for particle in xrange(start, state.obj.N):
        if not state.isactive(particle):
            continue

        if not quiet:
            print particle

        sys.stdout.flush()

        blocks = [state.blocks_particle(particle)[-1]]
        sample_state(state, blocks, stepout=stepout)

    return state.state.copy()

def sample_block(state, blockname, explode=True, stepout=0.1, quiet=False):
    if not quiet:
        print '{:-^39}'.format(' '+blockname.upper()+' ')

    blocks = [state.create_block(blockname)]

    if explode:
        blocks = state.explode(blocks[0])

    return sample_state(state, blocks, stepout)

def sample_block_list(state, blocklist, stepout=0.1):
    for bl in blocklist:
        sample_block(state, bl, stepout=stepout)
    return state.state.copy(), state.loglikelihood()

def do_samples(s, sweeps, burn, stepout=0.1, save_period=-1,
        prefix='cbamf', save_name=None, sigma=True, pos=True, quiet=False):
    h = []
    ll = []
    if not save_name:
        with tempfile.NamedTemporaryFile(suffix='.cbamf-state.pkl', prefix=prefix) as f:
            save_name = f.name

    for i in xrange(sweeps):
        if save_period > 0 and i % save_period == 0:
            with open(save_name, 'w') as tfile:
                pickle.dump([s,h,ll], tfile)

        if not quiet:
            print '{:=^79}'.format(' Sweep '+str(i)+' ')

        #sample_particles(s, stepout=stepout)
        if pos:
            sample_particle_pos(s, stepout=stepout, quiet=quiet)
        sample_particle_rad(s, stepout=stepout, quiet=quiet)
        sample_block(s, 'psf', stepout=stepout, quiet=quiet)
        sample_block(s, 'ilm', stepout=stepout, quiet=quiet)
        sample_block(s, 'off', stepout=stepout, quiet=quiet)
        sample_block(s, 'zscale', stepout=stepout, quiet=quiet)

        if sigma:
            sample_block(s, 'sigma', stepout=0.005, quiet=quiet)

        if i >= burn:
            h.append(s.state.copy())
            ll.append(s.loglikelihood())

    if save_period > 0 and save_name:
        os.remove(save_name)

    h = np.array(h)
    ll = np.array(ll)
    return h, ll

#=============================================================================
# Optimization methods like gradient descent
#=============================================================================
def diag_grad(state, blocks):<commit_msg>adding slab to sampler in runner<commit_after>import os
import sys
import numpy as np
import scipy.ndimage as nd
import tempfile
import pickle

from cbamf import const
from cbamf.mc import samplers, engines, observers

# Linear fit function, because I can
def linear_fit(x, y, sigma=1, N=100, burn=1000):
    from cbamf.states import LinearFit
    poly = np.polyfit(x,y,1)

    s = LinearFit(x, y, sigma=sigma)
    s.state[:2] = poly

    bl = s.explode(s.block_all())
    h = sample_state(s, s.explode(s.block_all()), N=burn, doprint=True, procedure='uniform')
    h = sample_state(s, s.explode(s.block_all()), N=burn, doprint=True, procedure='uniform')

    return s, h.get_histogram()

#=============================================================================
# Sampling methods that run through blocks and sample
#=============================================================================
def sample_state(state, blocks, stepout=1, slicing=True, N=1, doprint=False, procedure='uniform'):
    eng = engines.SequentialBlockEngine(state)
    opsay = observers.Printer()
    ohist = observers.HistogramObserver(block=np.array(blocks).any(axis=0))
    eng.add_samplers([samplers.SliceSampler1D(stepout, block=b, procedure=procedure) for b in blocks])

    eng.add_likelihood_observers(opsay) if doprint else None
    eng.add_state_observers(ohist)

    eng.dosteps(N)
    return ohist

def scan_ll(state, element, size=0.1, N=1000):
    start = state.state[element]

    ll = []
    vals = np.linspace(start-size, start+size, N)
    for val in vals:
        state.update(element, val)
        l = state.loglikelihood()
        ll.append(l)

    state.update(element, start)
    return vals, np.array(ll)

def scan_noise(image, state, element, size=0.01, N=1000):
    start = state.state[element]

    xs, ys = [], []
    for i in xrange(N):
        print i
        test = image + np.random.normal(0, state.sigma, image.shape)
        x,y = sample_ll(test, state, element, size=size, N=300)
        state.update(element, start)
        xs.append(x)
        ys.append(y)

    return xs, ys

def sample_particles(state, stepout=1, start=0):
    print '{:-^39}'.format(' POS / RAD ')
    for particle in xrange(start, state.obj.N):
        if not state.isactive(particle):
            continue

        print particle
        sys.stdout.flush()

        blocks = state.blocks_particle(particle)
        sample_state(state, blocks, stepout=stepout)

    return state.state.copy()

def sample_particle_pos(state, stepout=1, start=0, quiet=False):
    if not quiet:
        print '{:-^39}'.format(' POS ')

    for particle in xrange(start, state.obj.N):
        if not state.isactive(particle):
            continue

        if not quiet:
            print particle
        sys.stdout.flush()

        blocks = state.blocks_particle(particle)[:-1]
        sample_state(state, blocks, stepout=stepout)

    return state.state.copy()

def sample_particle_rad(state, stepout=1, start=0, quiet=False):
    if not quiet:
        print '{:-^39}'.format(' RAD ')

    for particle in xrange(start, state.obj.N):
        if not state.isactive(particle):
            continue

        if not quiet:
            print particle

        sys.stdout.flush()

        blocks = [state.blocks_particle(particle)[-1]]
        sample_state(state, blocks, stepout=stepout)

    return state.state.copy()

def sample_block(state, blockname, explode=True, stepout=0.1, quiet=False):
    if not quiet:
        print '{:-^39}'.format(' '+blockname.upper()+' ')

    blocks = [state.create_block(blockname)]

    if explode:
        blocks = state.explode(blocks[0])

    return sample_state(state, blocks, stepout)

def sample_block_list(state, blocklist, stepout=0.1):
    for bl in blocklist:
        sample_block(state, bl, stepout=stepout)
    return state.state.copy(), state.loglikelihood()

def do_samples(s, sweeps, burn, stepout=0.1, save_period=-1,
        prefix='cbamf', save_name=None, sigma=True, pos=True, quiet=False):
    h = []
    ll = []
    if not save_name:
        with tempfile.NamedTemporaryFile(suffix='.cbamf-state.pkl', prefix=prefix) as f:
            save_name = f.name

    for i in xrange(sweeps):
        if save_period > 0 and i % save_period == 0:
            with open(save_name, 'w') as tfile:
                pickle.dump([s,h,ll], tfile)

        if not quiet:
            print '{:=^79}'.format(' Sweep '+str(i)+' ')

        #sample_particles(s, stepout=stepout)
        if pos:
            sample_particle_pos(s, stepout=stepout, quiet=quiet)
        sample_particle_rad(s, stepout=stepout, quiet=quiet)
        sample_block(s, 'psf', stepout=stepout, quiet=quiet)
        sample_block(s, 'ilm', stepout=stepout, quiet=quiet)
        sample_block(s, 'off', stepout=stepout, quiet=quiet)
        sample_block(s, 'zscale', stepout=stepout, quiet=quiet)

        if s.slab:
            sample_block(s, 'slab', stepout=stepout, quiet=quiet)

        if sigma:
            sample_block(s, 'sigma', stepout=0.005, quiet=quiet)

        if i >= burn:
            h.append(s.state.copy())
            ll.append(s.loglikelihood())

    if save_period > 0 and save_name:
        os.remove(save_name)

    h = np.array(h)
    ll = np.array(ll)
    return h, ll

#=============================================================================
# Optimization methods like gradient descent
#=============================================================================
def diag_grad(state, blocks):
{'fraction_non_alphanumeric': 0.17772272685567483, 'fraction_numerical': 0.011744196715294981, 'mean_word_length': 3.157131960335622, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 9, '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': '1251095', 'n_tokens_mistral': 3660, 'n_tokens_neox': 3414, 'n_words': 810}