Fix animations (#3)

* allowed dyads to be colored

* added sample animation

* updated animation notebook

* deleted tmax

* comment out ffmpeg

.gitignore

@@ -139,4 +139,5 @@ dmypy.json
 cython_debug/
 
 Figures/
-test.*
+test.*
+.DS_Store

hypercontagion/visualization/animation.py

@@ -8,7 +8,7 @@
 
 
 def contagion_animation(
-    fig, H, transition_events, pos, node_colors, edge_colors, dt=1, fps=1
+    fig, H, transition_events, pos, node_colors, edge_colors, dt=1, fps=1, **args
 ):
     """Generate an animation object of contagion process.
 
@@ -30,6 +30,8 @@ def contagion_animation(
         the timestep at which to take snapshots of the system states
     fps : int, default: 1
         frames per second in the animation.
+    **args :
+        optional xgi draw args
 
     Returns
     -------
@@ -58,11 +60,18 @@ def contagion_animation(
             node_state[target] = status
 
         node_fc = {n: node_colors[node_state[n]] for n in H.nodes}
-        edge_fc = {e: edge_colors[edge_state[e]] for e in H.edges}
+        dyad_color = {
+            e: edge_colors[edge_state[e]] for e in H.edges.filterby("size", 2)
+        }
+        edge_fc = {
+            e: edge_colors[edge_state[e]] for e in H.edges.filterby("size", 2, "geq")
+        }
 
         # draw hypergraph
-        plt.title(f"Time is {interval_time}")
-        xgi.draw(H, pos=pos, node_fc=node_fc, edge_fc=edge_fc)
+        xgi.draw(
+            H, pos=pos, node_fc=node_fc, edge_fc=edge_fc, dyad_color=dyad_color, **args
+        )
+        plt.tight_layout()
         camera.snap()
 
     return camera.animate(interval=1000 / fps)

tutorials/tutorial_3_animations.ipynb

@@ -13,7 +13,8 @@
     "import numpy as np\n",
     "import random\n",
     "import networkx as nx\n",
-    "from IPython.display import HTML"
+    "from IPython.display import HTML\n",
+    "import matplotlib.animation as animation"
    ]
   },
   {
@@ -22,12 +23,17 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "n = 100\n",
-    "is_connected = False\n",
-    "while not is_connected:\n",
-    "    H = xgi.random_hypergraph(n, [0.03, 0.0002, 0.00001])\n",
-    "    is_connected = xgi.is_connected(H)\n",
-    "pos = xgi.barycenter_spring_layout(H)"
+    "H = xgi.load_xgi_data(\"diseasome\")\n",
+    "H.cleanup()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "xgi.unique_edge_sizes(H)"
    ]
   },
   {
@@ -37,10 +43,8 @@
    "outputs": [],
    "source": [
     "gamma = 0.05\n",
-    "beta2 = 0.1\n",
-    "beta3 = 0.05\n",
-    "beta4 = 0.01\n",
-    "tau = {1: 0, 2: beta2, 3: beta3, 4: beta4}\n",
+    "beta = 0.05\n",
+    "tau = {i: beta for i in xgi.unique_edge_sizes(H)}\n",
     "rho = 0.1"
    ]
   },
@@ -51,7 +55,7 @@
    "outputs": [],
    "source": [
     "transition_events = hc.discrete_SIR(\n",
-    "    H, tau, gamma, tmin=0, tmax=40, dt=1, rho=rho, return_event_data=True\n",
+    "    H, tau, gamma, tmin=0, dt=1, rho=rho, return_event_data=True\n",
     ")"
    ]
   },
@@ -70,22 +74,29 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "node_colors = {\"S\": \"green\", \"I\": \"red\", \"R\": \"blue\"}\n",
-    "edge_colors = {\"S\": \"green\", \"I\": \"red\", \"R\": \"blue\", \"OFF\": \"grey\"}\n",
+    "S_color = \"white\"\n",
+    "I_color = 'firebrick'\n",
+    "R_color = 'steelblue'\n",
+    "\n",
+    "node_colors = {\"S\": S_color, \"I\": I_color, \"R\": R_color}\n",
+    "edge_colors = {\"S\": S_color, \"I\": I_color, \"R\": R_color, \"OFF\": \"grey\"}\n",
     "fps = 1\n",
     "fig = plt.figure(figsize=(10, 10))\n",
-    "animation = hc.contagion_animation(\n",
-    "    fig, H, transition_events, pos, node_colors, edge_colors, fps=fps\n",
+    "anim = hc.contagion_animation(\n",
+    "    fig, H, transition_events, pos, node_colors, edge_colors, fps=fps, node_size=10, dyad_lw=1.5\n",
     ")\n",
-    "HTML(animation.to_jshtml())"
+    "# FFwriter = animation.FFMpegWriter(fps=fps)\n",
+    "# anim.save('animation.mp4', writer = FFwriter)"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": []
+   "source": [
+    "HTML(anim.to_jshtml())"
+   ]
   }
  ],
  "metadata": {
@@ -104,7 +115,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.4"
+   "version": "3.10.0"
   },
   "orig_nbformat": 4,
   "vscode": {