<commit_before>import itertools
import numpy as np

from pyqpbo import alpha_expansion_grid
from pyqpbo import alpha_expansion_graph
from daimrf import mrf
from lp_new import solve_lp


from IPython.core.debugger import Tracer
tracer = Tracer()


class StructuredProblem(object):
    def __init__(self):
        self.size_psi = None
        self.inference_calls = 0

    def psi(self, x, y):
        pass

    def inference(self, x, w):
        pass

    def loss(self, y, y_hat):
        # hamming loss:
        return np.sum(y != y_hat)

    def loss_augmented_inference(self, x, y, w):
        print("FALLBACK no loss augmented inference found")
        return self.inference(x, w)


class GridCRF(StructuredProblem):
    def __init__(self, n_states=2, inference_method='qpbo'):
        super(GridCRF, self).__init__()
        self.n_states = n_states
        self.inference_method = inference_method
        # n_states unary parameters, upper triangular for pairwise
        self.size_psi = n_states + n_states * (n_states + 1) / 2

    def psi(self, x, y):
        # x is unaries
        # y is a labeling
        ## y can also be continuous (from lp)
        if x.shape == y.shape:
            x_flat = x.reshape(-1, x.shape[-1])
            y_flat = y.reshape(-1, y.shape[-1])
            unaries_acc = np.sum(x_flat * y_flat, axis=0)
            labels = y
        else:
            ## unary features:
            gx, gy = np.ogrid[:x.shape[0], :x.shape[1]]
            selected_unaries = x[gx, gy, y]
            unaries_acc = np.bincount(y.ravel(), selected_unaries.ravel(),
                                      minlength=self.n_states)

            ##accumulated pairwise
            #make one hot encoding
            labels = np.zeros((y.shape[0], y.shape[1], self.n_states),
                              dtype=np.int)
            labels[gx, gy, y] = 1

        # vertical edges
        vert = np.dot(labels[1:, :, :].reshape(-1, self.n_states).T,
                      labels[:-1, :, :].reshape(-1, self.n_states))
        # horizontal edges
        horz = np.dot(labels[:, 1:, :].reshape(-1, self.n_states).T,
                      labels[:, :-1, :].reshape(-1, self.n_states))
        pw = vert + horz
        pw = pw + pw.T - np.diag(np.diag(pw))
        feature = np.hstack([unaries_acc,
                             pw[np.tri(self.n_states, dtype=np.bool)]])
        return feature

    def inference(self, x, w, relaxed=False):
        self.inference_calls += 1
        if w.shape != (self.size_psi,):
            raise ValueError("Got w of wrong shape. Expected %s, got %s" %
                             (self.size_psi, w.shape))
        unary_params = w[:self.n_states]
        pairwise_flat = np.asarray(w[self.n_states:])
        pairwise_params = np.zeros((self.n_states, self.n_states))
        pairwise_params[np.tri(self.n_states, dtype=np.bool)] = pairwise_flat
        pairwise_params = (pairwise_params + pairwise_params.T
                           - np.diag(np.diag(pairwise_params)))
        if self.inference_method == "qpbo":
            return self._inference_qpbo(x, unary_params, pairwise_params)
        elif  self.inference_method == "dai":
            return self._inference_dai(x, unary_params, pairwise_params)
        elif  self.inference_method == "lp":
            return self._inference_lp(x, unary_params, pairwise_params,
                                      relaxed)
        else:
            raise ValueError("inference_method must be 'qpbo' or 'dai', got %s"
                             % self.inference_method)

    def _inference_qpbo(self, x, unary_params, pairwise_params):
        unaries = (-1000 * unary_params * x).astype(np.int32)
        pairwise = (-1000 * pairwise_params).astype(np.int32)
        y = alpha_expansion_grid(unaries, pairwise)
        return y

    def _inference_dai(self, x, unary_params, pairwise_params):
        ## build graph
        inds = np.arange(x.shape[0] * x.shape[1]).reshape(x.shape[:2])
        inds = inds.astype(np.int64)
        horz = np.c_[inds[:, :-1].ravel(), inds[:, 1:].ravel()]
        vert = np.c_[inds[:-1, :].ravel(), inds[1:, :].ravel()]
        edges = np.vstack([horz, vert])
        log_unaries = unary_params * x.reshape(-1, self.n_states)
        max_entry = max(np.max(log_unaries), 1)
        unaries = np.exp(log_unaries / max_entry)

        y = mrf(unaries, edges, np.exp(pairwise_params / max_entry), alg='jt')
        y = y.reshape(x.shape[0], x.shape[1])

        return y

    def _inference_lp(self, x, unary_params, pairwise_params, relaxed=False):
        ## build graph
        inds = np.arange(x.shape[0] * x.shape[1]).reshape(x.shape[:2])
        inds = inds.astype(np.int64)
        horz = np.c_[inds[:, :-1].ravel(), inds[:, 1:].ravel()]
        vert = np.c_[inds[:-1, :].ravel(), inds[1:, :].ravel()]
        edges = np.vstack([horz, vert])
        unaries = unary_params * x.reshape(-1, self.n_states)
        y = solve_lp(-unaries, edges, -pairwise_params)
        n_fractional = np.sum(y.max(axis=-1) < .9)
        if n_fractional:
            print("fractional solutions found: %d" % n_fractional)
        if relaxed:
            return y.reshape(x.shape)
        y = np.argmax(y, axis=-1)
        y = y.reshape(x.shape[0], x.shape[1])
        return y
<commit_msg>Solve lp exact<commit_after>import itertools
import numpy as np

from pyqpbo import alpha_expansion_grid
from pyqpbo import alpha_expansion_graph
from daimrf import mrf
from lp_new import solve_lp


from IPython.core.debugger import Tracer
tracer = Tracer()


class StructuredProblem(object):
    def __init__(self):
        self.size_psi = None
        self.inference_calls = 0

    def psi(self, x, y):
        pass

    def inference(self, x, w):
        pass

    def loss(self, y, y_hat):
        # hamming loss:
        return np.sum(y != y_hat)

    def loss_augmented_inference(self, x, y, w):
        print("FALLBACK no loss augmented inference found")
        return self.inference(x, w)


class GridCRF(StructuredProblem):
    def __init__(self, n_states=2, inference_method='qpbo'):
        super(GridCRF, self).__init__()
        self.n_states = n_states
        self.inference_method = inference_method
        # n_states unary parameters, upper triangular for pairwise
        self.size_psi = n_states + n_states * (n_states + 1) / 2

    def psi(self, x, y):
        # x is unaries
        # y is a labeling
        ## y can also be continuous (from lp)
        if x.shape == y.shape:
            x_flat = x.reshape(-1, x.shape[-1])
            y_flat = y.reshape(-1, y.shape[-1])
            unaries_acc = np.sum(x_flat * y_flat, axis=0)
            labels = y
        else:
            ## unary features:
            gx, gy = np.ogrid[:x.shape[0], :x.shape[1]]
            selected_unaries = x[gx, gy, y]
            unaries_acc = np.bincount(y.ravel(), selected_unaries.ravel(),
                                      minlength=self.n_states)

            ##accumulated pairwise
            #make one hot encoding
            labels = np.zeros((y.shape[0], y.shape[1], self.n_states),
                              dtype=np.int)
            labels[gx, gy, y] = 1

        # vertical edges
        vert = np.dot(labels[1:, :, :].reshape(-1, self.n_states).T,
                      labels[:-1, :, :].reshape(-1, self.n_states))
        # horizontal edges
        horz = np.dot(labels[:, 1:, :].reshape(-1, self.n_states).T,
                      labels[:, :-1, :].reshape(-1, self.n_states))
        pw = vert + horz
        pw = pw + pw.T - np.diag(np.diag(pw))
        feature = np.hstack([unaries_acc,
                             pw[np.tri(self.n_states, dtype=np.bool)]])
        return feature

    def inference(self, x, w, relaxed=False):
        self.inference_calls += 1
        if w.shape != (self.size_psi,):
            raise ValueError("Got w of wrong shape. Expected %s, got %s" %
                             (self.size_psi, w.shape))
        unary_params = w[:self.n_states]
        pairwise_flat = np.asarray(w[self.n_states:])
        pairwise_params = np.zeros((self.n_states, self.n_states))
        pairwise_params[np.tri(self.n_states, dtype=np.bool)] = pairwise_flat
        pairwise_params = (pairwise_params + pairwise_params.T
                           - np.diag(np.diag(pairwise_params)))
        if self.inference_method == "qpbo":
            return self._inference_qpbo(x, unary_params, pairwise_params)
        elif  self.inference_method == "dai":
            return self._inference_dai(x, unary_params, pairwise_params)
        elif  self.inference_method == "lp":
            return self._inference_lp(x, unary_params, pairwise_params,
                                      relaxed)
        else:
            raise ValueError("inference_method must be 'qpbo' or 'dai', got %s"
                             % self.inference_method)

    def _inference_qpbo(self, x, unary_params, pairwise_params):
        unaries = (-1000 * unary_params * x).astype(np.int32)
        pairwise = (-1000 * pairwise_params).astype(np.int32)
        y = alpha_expansion_grid(unaries, pairwise)
        return y

    def _inference_dai(self, x, unary_params, pairwise_params):
        ## build graph
        inds = np.arange(x.shape[0] * x.shape[1]).reshape(x.shape[:2])
        inds = inds.astype(np.int64)
        horz = np.c_[inds[:, :-1].ravel(), inds[:, 1:].ravel()]
        vert = np.c_[inds[:-1, :].ravel(), inds[1:, :].ravel()]
        edges = np.vstack([horz, vert])
        log_unaries = unary_params * x.reshape(-1, self.n_states)
        max_entry = max(np.max(log_unaries), 1)
        unaries = np.exp(log_unaries / max_entry)

        y = mrf(unaries, edges, np.exp(pairwise_params / max_entry), alg='jt')
        y = y.reshape(x.shape[0], x.shape[1])

        return y

    def _inference_lp(self, x, unary_params, pairwise_params, relaxed=False):
        ## build graph
        inds = np.arange(x.shape[0] * x.shape[1]).reshape(x.shape[:2])
        inds = inds.astype(np.int64)
        horz = np.c_[inds[:, :-1].ravel(), inds[:, 1:].ravel()]
        vert = np.c_[inds[:-1, :].ravel(), inds[1:, :].ravel()]
        edges = np.vstack([horz, vert])
        unaries = unary_params * x.reshape(-1, self.n_states)
        y = solve_lp(-unaries, edges, -pairwise_params, exact=False)
        n_fractional = np.sum(y.max(axis=-1) < .9)
        #if n_fractional:
            #print("got fractional solution. trying again, this time exactly")
            #y = solve_lp(-unaries, edges, -pairwise_params, exact=True)
            #n_fractional = np.sum(y.max(axis=-1) < .9)
        if n_fractional:
            print("fractional solutions found: %d" % n_fractional)
        if relaxed:
            return y.reshape(x.shape)
        y = np.argmax(y, axis=-1)
        y = y.reshape(x.shape[0], x.shape[1])
        return y

{'fraction_non_alphanumeric': 0.1334008469895047, 'fraction_numerical': 0.0117842018044559, 'mean_word_length': 2.1019417475728157, 'pattern_counts': {'":': 6, '<': 6, '<?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': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5332191', 'n_tokens_mistral': 3774, 'n_tokens_neox': 3641, 'n_words': 777}