From 3a734a6759b94717b68f868f785d8502a1f2ea02 Mon Sep 17 00:00:00 2001
From: Jonathan Bloedow <jhb4@cornell.edu>
Date: Tue, 18 Jun 2024 15:54:52 -0700
Subject: [PATCH] Adding folium script to create animations from
 simulation_output.csv.

---
 jb/utils/folium_animate_from_sqlite.py | 119 +++++++++++++++++++++++++
 1 file changed, 119 insertions(+)
 create mode 100644 jb/utils/folium_animate_from_sqlite.py

diff --git a/jb/utils/folium_animate_from_sqlite.py b/jb/utils/folium_animate_from_sqlite.py
new file mode 100644
index 0000000..3afbee1
--- /dev/null
+++ b/jb/utils/folium_animate_from_sqlite.py
@@ -0,0 +1,119 @@
+import csv
+import sqlite3
+import random
+import numpy as np
+import sys
+import argparse
+
+import folium
+from folium.plugins import HeatMapWithTime
+
+# Connect to the SQLite database (or create it)
+conn = sqlite3.connect(':memory:')
+cursor = conn.cursor()
+
+def preproc( sim_report_file ):
+    # Create tables and import data from CSV files
+    cursor.execute("DROP TABLE IF EXISTS engwal")
+    cursor.execute("DROP TABLE IF EXISTS cities")
+
+    # Import simulation_output.csv into engwal table
+    with open(sim_report_file, 'r') as file:
+        reader = csv.reader(file)
+        headers = next(reader)
+        cursor.execute(f"CREATE TABLE engwal ({', '.join(headers)})")
+        cursor.executemany(f"INSERT INTO engwal VALUES ({', '.join(['?' for _ in headers])})", reader)
+
+    # Import cities.csv into cities table
+    with open('cities.csv', 'r') as file:
+        reader = csv.reader(file)
+        headers = next(reader)
+        cursor.execute(f"CREATE TABLE cities ({', '.join(headers)})")
+        cursor.executemany(f"INSERT INTO cities VALUES ({', '.join(['?' for _ in headers])})", reader)
+
+    # Create the view
+    cursor.execute("""
+    CREATE VIEW cases AS
+    SELECT Timestep, Node, Name, Latitude, Longitude, "New Infections"
+    FROM engwal, cities
+    WHERE engwal.Node = cities.ID
+    """)
+
+    conn.commit()
+
+
+def process( output_file ):
+    # Create a map centered at a specific location
+    birmingham_location = (52.485,-1.86)
+    m = folium.Map(location=(birmingham_location[0],birmingham_location[1]), zoom_start=8) # Create a list to store the data for HeatMapWithTime
+
+    start_time=800
+    cursor.execute(f'SELECT CAST(Timestep AS INT), Latitude, Longitude, CAST("New Infections" AS INT) FROM cases WHERE (CAST(Timestep AS INT)>{start_time})')
+    raw_data = cursor.fetchall()
+
+    # Group the data by timestep
+    grouped_data = {}
+    for row in raw_data:
+        timestep, lat, lon, new_infections = row
+        if timestep not in grouped_data:
+            grouped_data[timestep] = []
+        grouped_data[timestep].append([lat, lon, new_infections])
+
+    data = []
+    grouped_data = list(grouped_data.values())
+    locations=[[ float(elem[0]), float(elem[1]) ] for elem in grouped_data[0]]
+
+    max_cases = 500 # updating this manually
+    for t in range(len(grouped_data)-1):
+        this_row = []
+        for i, location in enumerate( locations ):
+            try:
+                def renorm( value ):
+                    if value > max_cases:
+                        value = max_cases
+                    offset = 1
+                    log_value = np.log(value + offset)
+                    log_min = np.log(0 + offset)
+                    log_max = np.log(max_cases + offset)
+                    return (log_value - log_min) / (log_max - log_min)
+                case_value = grouped_data[t][i][2]
+                case_value = renorm( case_value )
+                # random jitter seems to be essential!!! Yes, wack.
+                this_row.append( [ location[0], location[1], case_value+random.random()/2000 ] )
+            except Exception as ex:
+                print( f"Exception with t={t}, i={i}." )
+                raise ValueError( str( ex ) )
+        data.append( this_row )
+
+    heat_map = HeatMapWithTime(
+            data,
+            radius=15, 
+            gradient={0.0: 'blue', 0.2: 'green', 0.4: 'yellow', 0.6: 'orange', 0.8: 'red'},
+            min_opacity = 0.0,
+            max_opacity = 0.8,
+            #use_local_extrema=True,
+            auto_play=True 
+        )
+    heat_map.add_to(m)
+
+
+    # Save the map as an HTML file
+    m.save( output_file )
+
+if __name__ == "__main__":
+    # Set up argument parser
+    parser = argparse.ArgumentParser(description='Process simulation data.')
+    parser.add_argument('-i', '--input', type=str, default='simulation_output.csv',
+                        help='Input file name (default: simulation_output.csv)')
+    parser.add_argument('-o', '--output', type=str, default='sim_animation.html',
+                        help='Output file name (default: sim_animation.html)')
+    
+    # Parse command-line arguments
+    args = parser.parse_args()
+    
+    # Call the preproc function with the input file
+    preproc(args.input)
+    
+    # Call the process function with the output file
+    process(args.output)
+