<jupyter_start><jupyter_text><jupyter_code>!pip install --progress-bar off --quiet poetry
!pip install --progress-bar off --quiet git+https://github.com/oughtinc/ergo.git
!wget https://upload.wikimedia.org/wikipedia/commons/thumb/5/51/BlankMap-Equirectangular.svg/1280px-BlankMap-Equirectangular.svg.png
import ergo
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
from scipy import stats
from scipy.optimize import minimize
from math import pi

scale_factor = 1
k = 3

def gaussian_cross_entropy(param_list, distrib, log_distrib, param_num, distrib_len, m):
    result = 0
    total_mass = 1

    for i in range(param_num - 1):
      total_mass += np.exp(param_list[3*param_num + i])

    for x in range(distrib_len):
      curr_sum = 0
      for i in range(param_num):
        if i == 0:
          curr_sum += stats.norm.pdf(x, loc = param_list[3*i + m], scale = np.exp(param_list[3*i + 2]) * (m+1))/total_mass
        else:
          curr_sum += np.exp(param_list[3 * param_num + i - 1]) * stats.norm.pdf(x, loc = param_list[3*i + m], scale = np.exp(param_list[3*i + 2]) * (m+1))/total_mass
      result += -distrib[x] * (np.log(curr_sum) - log_distrib[x])

    return result

def objective(param_list, param_num):
    return scale_factor * (gaussian_cross_entropy(param_list, vertical_dist, log_vertical_dist, param_num, len(vertical_dist), 1) + gaussian_cross_entropy(param_list, horizontal_dist, log_horizontal_dist, param_num, len(horizontal_dist), 0))


metaculus = ergo.Metaculus()

height = 640
width = 1280

q_longitude = metaculus.get_question(3867, name="Longitude of WW3")
q_latitude = metaculus.get_question(3868, name="Latitude of WW3")

longitude_histogram = q_longitude.prediction_histogram
latitude_histogram = q_latitude.prediction_histogram

horizontal_dist = [x[2] for x in longitude_histogram]
horizontal_dist = horizontal_dist/np.sum(horizontal_dist)
log_horizontal_dist = np.log(horizontal_dist)

vertical_dist = [x[2] for x in latitude_histogram]
vertical_dist = vertical_dist/np.sum(vertical_dist)
log_vertical_dist = np.log(vertical_dist)



init_val = np.array(list(np.random.normal(loc = 100, scale = 10, size = 3*k)) + list(np.random.normal(loc = 0, scale = 0.1, size = k-1)))

for i in range(k):
  init_val[3*i + 2] = np.log(init_val[3*i + 2]/20)

res = minimize(objective, args = (k), method = "Powell", x0 = init_val, tol = 0.01)

print(res)
final_parameters = res.x

print(final_parameters)

heatmap_array = np.zeros((height, width))
sum_val = 0

total_mass = 1
for i in range(k - 1):
      total_mass += np.exp(final_parameters[3*k + i])

for i in range(int(width/10)):
  for j in range(int(height/10)):
    curr_val = 0
    for n in range(k):
      if n == 0:
        curr_val += stats.norm.pdf(10 * i/width * 200, loc = final_parameters[3*n], scale = np.exp(final_parameters[3*n + 2])) * stats.norm.pdf(10 * j/height * 200, loc = final_parameters[3*n + 1], scale = 2 * np.exp(final_parameters[3*n + 2]))/total_mass
      else:
        curr_val += np.exp(final_parameters[3*k + n - 1]) * stats.norm.pdf(10 * i/width * 200, loc = final_parameters[3*n], scale = np.exp(final_parameters[3*n + 2])) * stats.norm.pdf(10 * j/height * 200, loc = final_parameters[3*n + 1], scale = 2 * np.exp(final_parameters[3*n + 2]))/total_mass
    sum_val += 100 * curr_val
    for u in range(10):
      for v in range(10):
        heatmap_array[height - 10*j - v - 1][10*i + u] = curr_val

heatmap_array = heatmap_array/sum_val

world_map = Image.open('1280px-BlankMap-Equirectangular.svg.png')

plt.figure(figsize=(20,10))
plt.imshow(world_map)
plt.imshow(heatmap_array, cmap='hot', interpolation='bicubic', alpha = 0.5)
plt.show()<jupyter_output>[223.43506996  73.54831484   4.80766088 163.13204009 134.23091358
   2.49510539 121.29080639 135.84187147   2.33247892   2.94692307
   2.66567084]
{'fraction_non_alphanumeric': 0.1341308220921354, 'fraction_numerical': 0.06261135148892848, 'mean_word_length': 3.7636363636363637, 'pattern_counts': {'":': 0, '<': 4, '<?xml version=': 0, '>': 4, 'https://': 2, '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': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '209242', 'license': 'MIT', 'n_tokens_mistral': 1679, 'n_tokens_neox': 1465, 'n_words': 310, 'path': 'ww3_heatmap_improved.ipynb', 'repo_name': 'ege-erdil/ww3-heatmap-improved', 'size': '3929'}