Myo map heatmap (#519)

* heatmap draft * add data * improve heatmap * pre changing data to signals * transfer data to signals * visualize map * add modified map schema * html option * update heatmap * Renaming key variables to be more semantically clear (lat/lon vs. x/y) * Changing map+heatmap interactive output handling to write at user-specified location * add test for heatmap Co-authored-by: Joseph Cottam <[email protected]>
ciemss · Apr 1, 2024 · 3f16962 · 3f16962
1 parent a7da30f
commit 3f16962
Show file tree

Hide file tree

Showing 9 changed files with 48,804 additions and 7 deletions.
diff --git a/pyciemss/visuals/data/country_names.csv b/pyciemss/visuals/data/country_names.csv
diff --git a/pyciemss/visuals/data/world-110m.json b/pyciemss/visuals/data/world-110m.json
diff --git a/pyciemss/visuals/histogram.py b/pyciemss/visuals/histogram.py
@@ -1,11 +1,15 @@
+import json
 from numbers import Number
+from pathlib import Path
 from typing import Callable, List, Literal, Optional, Tuple, Union, overload
 
 import numpy as np
 import pandas as pd
 
 from . import vega
 
+_output_root = Path(__file__).parent / "data"
+
 
 def sturges_bin(data):
     """Determine number of bin susing sturge's rule.
@@ -209,3 +213,34 @@ def _mesh_to_heatmap(mesh_data):
         )
 
     return schema
+
+
+def map_heatmap(mesh: pd.DataFrame = None) -> vega.VegaSchema:
+    """
+    mesh -- (Optional) pd.DataFrame with columns
+        lon_start, lon_end, lat_start, lat_end, count for each grid
+    """
+
+    schema = vega.load_schema("map_heatmap.vg.json")
+    mesh_array = mesh.to_json(orient="records")
+    # load heatmap data
+    schema["data"] = vega.replace_named_with(
+        schema["data"], "mesh", ["values"], json.loads(mesh_array)
+    )
+    #
+    # add in map topology data
+    world_path = _output_root / "world-110m.json"
+    f = open(world_path)
+    world_data = json.load(f)
+    schema["data"] = vega.replace_named_with(
+        schema["data"], "world", ["values"], world_data
+    )
+
+    # add in country names
+    country_names_path = _output_root / "country_names.csv"
+
+    name_data = pd.read_csv(country_names_path).to_json(orient="records")
+    schema["data"] = vega.replace_named_with(
+        schema["data"], "names", ["values"], json.loads(name_data)
+    )
+    return schema
diff --git a/pyciemss/visuals/html/visualize_map.html b/pyciemss/visuals/html/visualize_map.html
@@ -0,0 +1,20 @@
+<!DOCTYPE html>
+<html>
+
+<head>
+  <script src="https://ajax.googleapis.com/ajax/libs/jquery/2.1.3/jquery.min.js"></script>
+  <script src="https://cdn.jsdelivr.net/npm/vega@5"></script>
+  <script src="https://cdn.jsdelivr.net/npm/[email protected]"></script>
+</head>
+
+<body>
+  <div id="vis" />
+  <script>
+    $.getJSON('modified_map_heatmap.json', function (spec) {
+
+      vegaEmbed("#vis", spec, { mode: "vega" }).then(console.log).catch(console.warn);
+    });
+  </script>
+</body>
+
+</html>
diff --git a/pyciemss/visuals/plots.py b/pyciemss/visuals/plots.py
@@ -1,13 +1,15 @@
 import json
-from typing import Any, Dict, Literal, Optional
+import shutil
+from pathlib import Path
+from typing import Any, Dict, Literal, Optional, Union
 
 import IPython.display
 import vl_convert
 
 from .barycenter import triangle_contour
 from .calibration import calibration
 from .graphs import attributed_graph, spring_force_graph
-from .histogram import heatmap_scatter, histogram_multi
+from .histogram import heatmap_scatter, histogram_multi, map_heatmap
 from .trajectories import trajectories
 from .vega import VegaSchema, orient_legend, pad, rescale, resize, set_title
 
@@ -25,21 +27,34 @@
     "attributed_graph",
     "spring_force_graph",
     "heatmap_scatter",
+    "map_heatmap",
 ]
 
 
-def save_schema(schema: Dict[str, Any], path: str):
+def save_schema(schema: Dict[str, Any], path: Path):
     """Save the schema using common convention"""
     with open(path, "w") as f:
         json.dump(schema, f, indent=3)
 
 
+def check_geoscale(schema):
+    geoscale = False
+    if "signals" in schema.keys():
+        for i in range(len(schema["signals"])):
+            signal = schema["signals"][i]
+            if "on" in signal.keys():
+                if "geoscale" in signal["on"][0]["update"].lower():
+                    geoscale = True
+    return geoscale
+
+
 def ipy_display(
     schema: Dict[str, Any],
     *,
     format: Literal["png", "svg", "PNG", "SVG", "interactive", "INTERACTIVE"] = "png",
     force_clear: bool = False,
     dpi: Optional[int] = None,
+    output_root: Union[str, Path, None] = None,
     **kwargs,
 ):
     """Wrap for dispaly in an ipython notebook.
@@ -49,6 +64,7 @@ def ipy_display(
               Format specifier is case-insensitive.
     force_clear -- Force clear the result cell (sometimes required after an error)
     dpi -- approximates DPI for output (other factors apply)
+    output_root -- Location of output files. String name of new folder will be converted to Pathlib Path
     **kwargs -- Passed on to the selected vl_convert function
 
     The vlc_convert PNG export function takes a 'scale' factor,
@@ -71,7 +87,28 @@ def ipy_display(
     if force_clear:
         IPython.display.clear_output(wait=True)
 
-    if format in ["interactive", "INTERACTIVE"]:
+    if isinstance(output_root, str):
+        output_root = Path(output_root)
+
+    if check_geoscale(schema):
+        if output_root is None:
+            raise ValueError(
+                "Must supply an writeable output directory when visualizing this type of schema"
+            )
+
+        output_schema = output_root / "modified_map_heatmap.json"
+        output_html = output_root / "visualize_map.html"
+        input_html = Path(__file__).parent / "html" / "visualize_map.html"
+
+        if not output_root.exists():
+            output_root.mkdir(parents=True)
+
+        shutil.copy(input_html, output_html)
+        save_schema(schema, output_schema)
+        print(
+            f"Schema includes 'geoscale' which can't be interactively rendered. View at {output_html}"
+        )
+    elif format in ["interactive", "INTERACTIVE"]:
         bundle = {"application/vnd.vega.v5+json": schema}
         print("", end=None)
         IPython.display.display(bundle, raw=True)