<commit_before>import sys
import os.path
import os

os.environ["THEANO_FLAGS"] = "floatX=float64"

from collections import defaultdict
import json
import theano
from theano import tensor as T
from theano.ifelse import ifelse
import glob
import progressbar
import numpy as np

def link_cost(v, winsor=10.0):
    a = (v ** 2) / 2
    b = winsor * (abs(v) - winsor / 2)
    return T.switch(T.lt(abs(v), winsor), a, b)

if __name__ == "__main__":

    imsize = 50

    R1 = T.scalar()
    R2 = T.scalar()
    T1 = T.col()
    T2 = T.col()
    pts1 = T.matrix()
    pts2 = T.matrix()

    rot1 = T.as_tensor([T.cos(R1), -T.sin(R1),
                        T.sin(R1), T.cos(R1)]).reshape((2, 2))
    rot2 = T.as_tensor([T.cos(R2), -T.sin(R2),
                        T.sin(R2), T.cos(R2)]).reshape((2, 2))
    new_pts1 = T.dot(rot1, pts1) + T1 * imsize
    new_pts2 = T.dot(rot2, pts2) + T2 * imsize
    ptdists = T.sqrt(T.sum((new_pts1 - new_pts2) ** 2, axis=0) + 0.01)
    cost = T.sum(link_cost(ptdists))
    grads = T.grad(cost, wrt=[R1, R2, T1, T2])

    costgrad = theano.function([R1, R2, T1, T2, pts1, pts2],
                               outputs=[cost] + grads)

    distsfun = theano.function([R1, R2, T1, T2, pts1, pts2],
                               outputs=ptdists)

    rot_angles = defaultdict(lambda: 0.0)
    translations = defaultdict(lambda: np.zeros((2, 1)))


    all_matches = {}
    match_files = glob.glob(os.path.join(sys.argv[1], '*sift_matches*000001*000001*.json'))
    pbar = progressbar.ProgressBar()
    for f in pbar(match_files):
        data = json.load(open(f))
        pts1 = np.array([c["p1"]["w"] for c in data[0]["correspondencePointPairs"]]).T
        pts2 = np.array([c["p2"]["w"] for c in data[0]["correspondencePointPairs"]]).T
        url1 = data[0]["url1"]
        url2 = data[0]["url2"]
        dists = distsfun(rot_angles[url1], rot_angles[url2],
                          translations[url1], translations[url2],
                          pts1, pts2)
        keep = dists < np.median(dists)
        shift = np.mean(pts1[:, keep], axis=1).reshape((2, 1))
        all_matches[url1, url2] = (pts1[:, keep] - shift, pts2[:, keep] - shift)

    stepsize = 1e-10

    prev_c = np.inf
    for iter in range(20000):
        c = 0
        dists = []
        for (url1, url2), (pts1, pts2) in all_matches.iteritems():
            cg = costgrad(rot_angles[url1], rot_angles[url2],
                          translations[url1], translations[url2],
                          pts1, pts2)
            d = distsfun(rot_angles[url1], rot_angles[url2],
                         translations[url1], translations[url2],
                         pts1, pts2)
            dists.append(np.median(d))
            cg = [v / pts1.shape[1] for v in cg]
            c += cg[0]
            g = cg[1:]
            rot_angles[url1] -= stepsize * g[0]
            rot_angles[url2] -= stepsize * g[1]
            translations[url1] -= stepsize * g[2]
            translations[url2] -= stepsize * g[3]
        print "cost", c, stepsize, np.median(dists)
        if c < prev_c:
            stepsize *= 1.05
        else:
            stepsize *= 0.5
        prev_c = c
<commit_msg>2d rigid tile alignment for MFOVs - needs more testing<commit_after>import sys
import os.path
import os

from collections import defaultdict
import json

import glob
import progressbar
import numpy as np

def dist(p1, p2):
    return np.sqrt(((p1 - p2) ** 2).sum(axis=0))

def find_rotation(p1, p2, scale):
    U, S, VT = np.linalg.svd(np.dot(p1, p2.T))
    print S
    R = np.dot(VT.T, U.T)
    angle = stepsize * np.arctan2(R[1, 0], R[0, 0])
    return np.array([[np.cos(angle), -np.sin(angle)],
                     [np.sin(angle),  np.cos(angle)]])

if __name__ == "__main__":
    # all matched pairs between point sets
    all_matches = {}
    # all points from a given tile
    all_pts = defaultdict(list)

    match_files = glob.glob(os.path.join(sys.argv[1], '*sift_matches*000004*000004*.json'))
    pbar = progressbar.ProgressBar()
    for f in pbar(match_files):
        data = json.load(open(f))
        # point arrays are 2xN
        pts1 = np.array([c["p1"]["w"] for c in data[0]["correspondencePointPairs"]]).T
        pts2 = np.array([c["p2"]["w"] for c in data[0]["correspondencePointPairs"]]).T
        url1 = data[0]["url1"]
        url2 = data[0]["url2"]
        if pts1.size > 0:
            all_matches[url1, url2] = (pts1, pts2)

    maxiter = 2000
    epsilon = 5
    stepsize = 0.01
    prev_meanmed = np.inf

    T = defaultdict(lambda: np.zeros((2, 1)))
    R = defaultdict(lambda: np.eye(2))
    for iter in range(maxiter):
        if stepsize < 1e-30:
            break

        # transform points by the current trans/rot
        trans_matches = {(k1, k2): (np.dot(R[k1], p1) + T[k1],
                                    np.dot(R[k2], p2) + T[k2])
                         for (k1, k2), (p1, p2) in all_matches.iteritems()}

        # mask off all points more than epsilon past the median
        distances = {k: dist(p1, p2) for k, (p1, p2) in trans_matches.iteritems()}
        masks = {k : d < (np.median(d) + epsilon) for k, d in distances.iteritems()}
        masked_matches = {k: (p1[:, masks[k]], p2[:, masks[k]]) for k, (p1, p2) in trans_matches.iteritems()}

        median_dists = [np.median(d) for d in distances.values()]
        medmed = np.median(median_dists)
        meanmed = np.mean(median_dists)
        print("med-med distance: {}, mean-med distance: {}  SZ: {}".format(medmed, meanmed, stepsize))
        if meanmed < prev_meanmed:
            stepsize *= 1.1
            if stepsize > 1:
                stepsize = 1
        else:
            stepsize *= 0.5
        prev_meanmed = meanmed

        # find points and their matches from other groups for each tile
        self_points = defaultdict(list)
        other_points = defaultdict(list)
        for (k1, k2), (p1, p2) in masked_matches.iteritems():
            self_points[k1].append(p1)
            self_points[k2].append(p2)
            other_points[k1].append(p2)
            other_points[k2].append(p1)
        self_points = {k : np.hstack(p) for k, p in self_points.iteritems()}
        other_points = {k : np.hstack(p) for k, p in other_points.iteritems()}

        # update translations
        self_centers = {k : np.mean(p, axis=1).reshape((2, 1)) for k, p in self_points.iteritems()}
        other_centers = {k : np.mean(p, axis=1).reshape((2, 1)) for k, p in other_points.iteritems()}

        # find best rotation, multiply the angle of rotation by a stepsize, and update the rotations
        new_R = {k: find_rotation(self_points[k] - self_centers[k],
                                  other_points[k] - other_centers[k],
                                  stepsize)
                 for k in self_centers}
        R = {k: np.dot(R[k], new_R[k]) for k in R}
        deltas = {k: stepsize * (other_centers[k] - np.dot(new_R[k], self_centers[k])) for k in self_centers}
        T = {k: old_translation + deltas[k] for k, old_translation in T.iteritems()}
{'fraction_non_alphanumeric': 0.14457660719342963, 'fraction_numerical': 0.0356839422259983, 'mean_word_length': 2.17438202247191, 'pattern_counts': {'":': 2, '<': 8, '<?xml version=': 0, '>': 5, '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': '3489379', 'n_tokens_mistral': 2633, 'n_tokens_neox': 2533, 'n_words': 603}