<jupyter_start><jupyter_text>Functions<jupyter_code>def dist(x1, y1, x2, y2):
    return np.sqrt(np.power(x1 - x2, 2) + np.power(y1 - y2, 2))

def ny_scatter2(ax, df, coord_col_name = "test_{}_rank"):
    # fig, ax = plt.subplots(figsize = (s, s))
    for ind, county in enumerate(df.county_name.unique()):
        county_df = df[df.county_name == county]
        ax.scatter(
            y = county_df[coord_col_name.format("lat")],
            x = county_df[coord_col_name.format("lng")],
            s = county_df.sqrt_mult * 2,
            marker = "s",
            label = county,
            alpha = 0.7,
            color = colors[ind % 5],
            edgecolor = "none")
    # ax.legend(loc = 9, bbox_to_anchor = (0.5,0))
    # return fig

def check_intersect_ver2(row):
    def rev_pair(l):
        return [l, l[::-1]]
    
    return (min(
        [(row["new_scaled_{}_rank_{}".format(i, j[0])] - row["sqrt_mult_{}".format(j[0])] / 2) <=
         (row["new_scaled_{}_rank_{}".format(i, j[1])] + row["sqrt_mult_{}".format(j[0])] / 2)
        for i in ["lng", "lat"] for j in rev_pair(["x", "y"])]))

def get_min_max(df, center_col_name = "prop_scaled_{}_rank", side_col_name = "sqrt_mult"):
    for si, i in enumerate(["min", "max"]):
        for j in ["lat", "lng"]:
            df[("{}_" + center_col_name).format(i, j)] = (
                df[center_col_name.format(j)] + (2 * si - 1) *
                (df[side_col_name] / 2))
    return df

def get_pairs(df, id_col_name, select_col_names):
    df["dummy"] = 1
    pairs = df[select_col_names + ["dummy"]].merge(right = df[select_col_names+ ["dummy"]], on = "dummy", how = "left")
    return pairs[~(pairs[id_col_name + "_x"] == pairs[id_col_name + "_y"])].drop(columns = ["dummy"]) 

def get_intersect(df, center_col_name = "new_scaled_{}_rank_{}"):
    col_pat = "{}_" + center_col_name
    return (
        (df[col_pat.format("min", "lat", "x")] <= df[col_pat.format("max", "lat", "y")]) &
        (df[col_pat.format("min", "lat", "y")] <= df[col_pat.format("max", "lat", "x")]) &
        (df[col_pat.format("min", "lng", "x")] <= df[col_pat.format("max", "lng", "y")]) &
        (df[col_pat.format("min", "lng", "y")] <= df[col_pat.format("max", "lng", "x")]))    

def compress(ny_map_county_shapes):
    # get square pairs
    select_cols = ["square_id"] + ["new_scaled_{}_rank".format(i) for i in ["lat", "lng"]]
    ny_map_county_square_pairs = get_pairs(ny_map_county_shapes, "square_id", select_cols)    
    
    # calculate each square's force on each other county
    # scaled by the distance between squares
    d = dist(*[ny_map_county_square_pairs["new_scaled_{}_rank_{}".format(j, i)] for i in ["x", "y"] for j in ["lat", "lng"]])
    min_d = d.min()
    max_d = d.max()
    a = - 0.001 / ((min_d - max_d) ** 2)
    d_factor = 1 / d # a * (min_d - d) * (max_d - d)
    
    for i in ["lat", "lng"]:
        ny_map_county_square_pairs["force_{}".format(i)] = (d_factor * 
            (ny_map_county_square_pairs["new_scaled_{}_rank_y".format(i)] - ny_map_county_square_pairs["new_scaled_{}_rank_x".format(i)]))
        
    # add all forces together
    sum_forces = ny_map_county_square_pairs.groupby(
        ["square_id_x"], as_index = False).agg(
        {"force_{}".format(i): "sum" for i in ["lat", "lng"]})
    
    # join sum of forces to squares data
    ny_map_county_shapes = ny_map_county_shapes.merge(
        right = sum_forces,
        left_on = ["square_id"],
        right_on = ["square_id_x"],
        how = "left")
    
    # calculate a proposed new center for each square 
    for i in ["lat", "lng"]:
        ny_map_county_shapes["prop_scaled_{}_rank".format(i)] = (ny_map_county_shapes["new_scaled_{}_rank".format(i)] +
            ny_map_county_shapes["force_{}".format(i)])
        
    # check intersection
    ny_map_county_shapes = get_min_max(ny_map_county_shapes)
    select_cols = (
        ["square_id", "sqrt_mult"] +
        ["{}prop_scaled_{}_rank".format(j, i) for i in ["lat", "lng"] for j in ["", "min_", "max_"]])
    ny_map_county_square_pairs = get_pairs(ny_map_county_shapes, "square_id", select_cols)
    ny_map_county_square_pairs["intersect"] = get_intersect(ny_map_county_square_pairs, center_col_name = "prop_scaled_{}_rank_{}")
    
    square_intersect = ny_map_county_square_pairs.groupby(["square_id_x"], as_index = False)["intersect"].sum()
    ny_map_county_shapes = ny_map_county_shapes.merge(right = square_intersect, on = "square_id_x", how = "left")
    
    # apply force if no intersect
    for i in ["lat", "lng"]:
        ny_map_county_shapes["new_scaled_{}_rank".format(i)] = ny_map_county_shapes.apply(
            lambda x: x["new_scaled_{}_rank".format(i)] if x["intersect"] else x["prop_scaled_{}_rank".format(i)], axis = 1)
    
    return [ny_map_county_shapes.drop(columns = ["force_{}".format(i) for i in ["lat", "lng"]] + ["square_id_x", "intersect"]), ny_map_county_square_pairs]<jupyter_output><empty_output><jupyter_text>Import data<jupyter_code>county_shapes = pd.read_csv("./ny_county_shapes.csv").set_index("county_name")
ny_map_county_shapes = pd.read_csv("./ny_map_county_squares.csv")<jupyter_output><empty_output><jupyter_text>Untangle<jupyter_code>fig, ax = plt.subplots(figsize = (5, 5))
ny_scatter2(ax, ny_map_county_shapes)
for i in ["lat", "lng"]:
    ny_map_county_shapes["new_scaled_{}_rank".format(i)] = ny_map_county_shapes["test_{}_rank".format(i)]
ny_map_county_shapes["square_id"] = ny_map_county_shapes.index
# check to make sure squares dont intersect
ny_map_county_shapes_min_max = get_min_max(ny_map_county_shapes, center_col_name = "new_scaled_{}_rank")
select_cols = ["square_id", "sqrt_mult"] + ["{}new_scaled_{}_rank".format(j, i) for i in ["lat", "lng"] for j in ["", "min_", "max_"]]
ny_map_county_square_pairs = get_pairs(ny_map_county_shapes, "square_id", select_cols)
ny_map_county_square_pairs["intersect"] = get_intersect(ny_map_county_square_pairs)
ny_map_county_square_pairs.intersect.min(), ny_map_county_square_pairs.intersect.max()<jupyter_output><empty_output><jupyter_text>Intersect test<jupyter_code># need a test df with new_scaled_{}_rank_{} (lat, lng, x, y) and sqrt_mult_{} (x, y)
intersect_test = [
    [0, 0, 1, 1],
    [1, 1, 1, 1],
    [2, 2, 2, 1],
    [3, 3, 2, 1],
    [5, 5, 5, 1],
    [10, 10, 10, 1],
]
intersect_test_df = pd.DataFrame(
    intersect_test,
    columns = ["square_id"] + ["new_scaled_{}_rank".format(i) for i in ["lat", "lng"]] + ["sqrt_mult"])

intersect_test_df = get_min_max(intersect_test_df, center_col_name = "new_scaled_{}_rank")
intersect_test_pairs_df = get_pairs(intersect_test_df, "square_id", list(intersect_test_df.columns))
intersect_test_pairs_df
get_intersect(intersect_test_pairs_df)
compress(intersect_test_df)[0]
sns.scatterplot(
    data = intersect_test_df,
    x = "new_scaled_lat_rank",
    y = "new_scaled_lng_rank",
    marker = "s",
    s = 100)
def apply_compress(df, ntimes):
    if ntimes == 0:
        return df
    else:
        return apply_compress(compress(df)[0], ntimes - 1)
sns.scatterplot(
    data = apply_compress(intersect_test_df, 11),
    x = "new_scaled_lat_rank",
    y = "new_scaled_lng_rank",
    marker = "s",
    s = 100)
sns.boxenplot(ny_map_county_shapes.sqrt_mult)
niter = 20
prev = ny_map_county_shapes
test_dfs = []

for i in tqdm(range(niter)):
    test_dfs.append(compress(prev)[0])
    # print(compress(prev)[1].intersect.min(), compress(prev)[1].intersect.max())
    prev = test_dfs[-1]
fig, ax = plt.subplots(figsize = (5, niter * 5), nrows = niter + 1)#, sharey = True, sharex = True)
ny_scatter2(ax[0], ny_map_county_shapes)
for i in range(niter):
    ny_scatter2(ax[i + 1], test_dfs[i], coord_col_name = "new_scaled_{}_rank")
test_dfs[-1].drop(columns = ["state", "nyc", 'min_lng_rank', 'min_lat_rank', 'max_lng_rank', 'max_lat_rank',
       'num_squares', 'county', 'ttl', 'pre_mult', 'mult']).head()
test = get_min_max(test_dfs[-1])
test[test.county == "Bronx"].drop(columns = ["state", "nyc", 'min_lng_rank', 'min_lat_rank', 'max_lng_rank', 'max_lat_rank',
       'num_squares', 'county', 'ttl', 'pre_mult', 'mult', 'mid_lng_rank', 'mid_lat_rank', 'scaled_lng_rank', 'scaled_lat_rank',
       'test_lat_rank', 'test_lng_rank'])
test.to_csv("./untangled_county_squares.csv", index = False)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.17198218798892767, 'fraction_numerical': 0.011433385485618005, 'mean_word_length': 4.213299874529485, 'pattern_counts': {'":': 0, '<': 22, '<?xml version=': 0, '>': 17, '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}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '309724', 'license': 'MIT', 'n_tokens_mistral': 3408, 'n_tokens_neox': 3327, 'n_words': 578, 'path': 'maps/UntanglingCountiesVer2.ipynb', 'repo_name': 'gautsi/gen-purpose-repo', 'size': '8309'}