Match infections per node (#10)

* matching ipn, upversion python

* test ipn matching

* Update contagion.py

* remove jit

* format with black

* organize

fitting_ipn.ipynb

@@ -0,0 +1,155 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "import matplotlib.pyplot as plt\n",
+    "from lcs import *\n",
+    "import networkx as nx"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "G = nx.karate_club_graph()\n",
+    "A = nx.adjacency_matrix(G, weight=None).todense()\n",
+    "n = A.shape[0]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "b0 = 0.5\n",
+    "a = 0.01\n",
+    "alpha = 1\n",
+    "max_iter = 100\n",
+    "tol = 1e-3\n",
+    "\n",
+    "cf1 = lambda nu, b: 1 - (1 - b) ** nu\n",
+    "cf2 = lambda nu, b: b * (nu >= 2)\n",
+    "\n",
+    "gamma = 0.1\n",
+    "b = 0.03\n",
+    "rho0 = 1\n",
+    "tmax = 1000\n",
+    "\n",
+    "realizations = 100\n",
+    "\n",
+    "x0 = np.zeros(n)\n",
+    "x0[list(random.sample(range(n), int(rho0 * n)))] = 1\n",
+    "c1 = cf1(np.arange(n), b)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# mean fitting\n",
+    "mode = \"mean\"\n",
+    "ipn_c1 = 0\n",
+    "for _ in range(realizations):\n",
+    "    x = contagion_process(A, gamma, c1, x0, tmin=0, tmax=tmax)\n",
+    "    ipn_c1 += infections_per_node(x, mode) / realizations\n",
+    "\n",
+    "f = lambda b: ipn_func(b, ipn_c1, cf2, gamma, A, rho0, 100, tmax, mode)\n",
+    "b1, bvec1, fvec1 = robbins_monro_solve(\n",
+    "    f, b0, a, alpha, max_iter, tol, verbose=True, loss=\"function\", return_values=True\n",
+    ")\n",
+    "\n",
+    "# median fitting\n",
+    "mode = \"median\"\n",
+    "ipn_c1 = 0\n",
+    "for _ in range(realizations):\n",
+    "    x = contagion_process(A, gamma, c1, x0, tmin=0, tmax=tmax)\n",
+    "    ipn_c1 += infections_per_node(x, mode) / realizations\n",
+    "\n",
+    "f = lambda b: ipn_func(b, ipn_c1, cf2, gamma, A, rho0, 100, tmax, mode)\n",
+    "b2, bvec2, fvec2 = robbins_monro_solve(\n",
+    "    f, b0, a, alpha, max_iter, tol, verbose=True, loss=\"function\", return_values=True\n",
+    ")\n",
+    "\n",
+    "# max fitting\n",
+    "mode = \"max\"\n",
+    "ipn_c1 = 0\n",
+    "for _ in range(realizations):\n",
+    "    x = contagion_process(A, gamma, c1, x0, tmin=0, tmax=tmax)\n",
+    "    ipn_c1 += infections_per_node(x, mode) / realizations\n",
+    "\n",
+    "f = lambda b: ipn_func(b, ipn_c1, cf2, gamma, A, rho0, 100, tmax, mode)\n",
+    "b3, bvec3, fvec3 = robbins_monro_solve(\n",
+    "    f, b0, a, alpha, max_iter, tol, verbose=True, loss=\"function\", return_values=True\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "plt.figure()\n",
+    "plt.subplot(121)\n",
+    "plt.plot(bvec1, label=\"mean\")\n",
+    "plt.plot(bvec2, label=\"median\")\n",
+    "plt.plot(bvec3, label=\"max\")\n",
+    "plt.xlabel(\"iterations\")\n",
+    "plt.ylabel(\"fitted probability\")\n",
+    "\n",
+    "plt.subplot(122)\n",
+    "plt.plot(fvec1, label=\"mean\")\n",
+    "plt.plot(fvec2, label=\"median\")\n",
+    "plt.plot(fvec3, label=\"max\")\n",
+    "plt.legend()\n",
+    "plt.xlabel(\"iterations\")\n",
+    "plt.ylabel(\"diff. between expected IPNs\")\n",
+    "plt.tight_layout()\n",
+    "\n",
+    "plt.savefig(\"test.png\", dpi=1000)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "hyper",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.0"
+  },
+  "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

lcs/contagion.py

@@ -40,6 +40,66 @@ def contagion_process(A, gamma, c, x0, tmin=0, tmax=20, dt=1, random_seed=None):
     if not isinstance(A, ndarray):
         A = A.todense()
 
+    if A.dtype != "int":
+        A = np.array(A, dtype=int)
+
+    if random_seed is not None:
+        np.random.seed(random_seed)
+
+    n, m = np.shape(A)
+    if n != m:
+        raise Exception("Matrix must be square!")
+
+    T = int((tmax - tmin) / dt)
+    x = np.zeros((T, n), dtype=int)
+    x[0] = x0.copy()
+
+    for t in range(T - 1):
+        u = np.random.random(n)
+        nu = A @ x[t]
+        x[t + 1] = (1 - x[t]) * (u <= c[nu] * dt) + x[t] * (u > gamma * dt)
+    return x
+
+
+def contagion_process_loop(A, gamma, c, x0, tmin=0, tmax=20, dt=1, random_seed=None):
+    """A neighborhood-based contagion process on pairwise networks
+
+    Parameters
+    ----------
+    A : Numpy ndarray
+        The adjacency matrix
+    gamma : float
+        The healing rate
+    c : Numpy ndarray
+        A 1d vector of the contagion rates. Should be N x 1.
+    x0 : Numpy ndarray
+        A 1D vector of the initial nodal states, either 0 or 1.
+        Should be N x 1.
+    tmin : int, optional
+        The time at which to start the simulation, by default 0
+    tmax : float, optional
+        The time at which to terminate the simulation, by default 20
+    dt : float, optional
+        The time step, by default 1
+    random_seed : int, optional
+        The seed for the random process, by default None
+
+    Returns
+    -------
+    Numpy ndarray
+        The time series of the contagion process. Should be T x N.
+
+    Raises
+    ------
+    Exception
+        If adjacency matrix isn't square
+    """
+    if not isinstance(A, ndarray):
+        A = A.todense()
+
+    if A.dtype != "int":
+        A = np.array(A, dtype=int)
+
     if random_seed is not None:
         random.seed(random_seed)
 
@@ -48,7 +108,7 @@ def contagion_process(A, gamma, c, x0, tmin=0, tmax=20, dt=1, random_seed=None):
         raise Exception("Matrix must be square!")
 
     T = int((tmax - tmin) / dt)
-    x = np.zeros((T, n))
+    x = np.zeros((T, n), dtype=int)
     x[0] = x0.copy()
 
     for t in range(T - 1):
@@ -59,7 +119,7 @@ def contagion_process(A, gamma, c, x0, tmin=0, tmax=20, dt=1, random_seed=None):
             if x[t, i] == 1 and random.random() <= gamma * dt:
                 x[t + 1, i] = 0
             elif x[t, i] == 0:
-                infected_nbrs = int(np.round(A[i] @ x[t]))
+                infected_nbrs = A[i] @ x[t]
                 if random.random() <= c[infected_nbrs] * dt:
                     x[t + 1, i] = 1
     return x