jinns 1.4.0__tar.gz → 1.5.0__tar.gz

{jinns-1.4.0 → jinns-1.5.0}/Notebooks/ODE/1D_Generalized_Lotka_Volterra.ipynb

@@ -133,7 +133,7 @@
     "    (eqx.nn.Linear, 64, 64),\n",
     "    (jax.nn.tanh,),\n",
     "    (eqx.nn.Linear, 64, 3),\n",
-    "    (jnp.exp,)\n",
+    "    (jnp.exp,),\n",
     ")\n",
     "key, subkey = random.split(key)"
    ]
@@ -154,19 +154,14 @@
    "outputs": [],
    "source": [
     "nt = 5000\n",
-    "method = 'uniform'\n",
+    "method = \"uniform\"\n",
     "tmin = 0\n",
     "tmax = 1\n",
     "\n",
     "Tmax = 30\n",
     "key, subkey = random.split(key)\n",
     "train_data = jinns.data.DataGeneratorODE(\n",
-    "    key=subkey,\n",
-    "    nt=nt,\n",
-    "    tmin=tmin,\n",
-    "    tmax=tmax,\n",
-    "    temporal_batch_size=512,\n",
-    "    method=method\n",
+    "    key=subkey, nt=nt, tmin=tmin, tmax=tmax, temporal_batch_size=512, method=method\n",
     ")"
    ]
   },
@@ -185,7 +180,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "u, init_nn_params = jinns.nn.PINN_MLP.create(key=subkey, eqx_list=eqx_list, eq_type=\"ODE\")"
+    "u, init_nn_params = jinns.nn.PINN_MLP.create(\n",
+    "    key=subkey, eqx_list=eqx_list, eq_type=\"ODE\"\n",
+    ")"
    ]
   },
   {
@@ -207,7 +204,8 @@
    "source": [
     "# initial conditions for each species\n",
     "import numpy as onp\n",
-    "N_0 = onp.array([10., 7., 4.])\n",
+    "\n",
+    "N_0 = onp.array([10.0, 7.0, 4.0])\n",
     "# growth rates for each species\n",
     "growth_rates = jnp.array([0.1, 0.5, 0.8])\n",
     "# carrying capacity for each species\n",
@@ -215,10 +213,11 @@
     "# interactions\n",
     "# NOTE that interaction with oneself is 0\n",
     "# NOTE minus sign\n",
-    "interactions = (-jnp.array([0, 0.001, 0.001]),\n",
-    "                -jnp.array([0.001, 0, 0.001]), \n",
-    "                -jnp.array([0.001, 0.001, 0])\n",
-    "                )"
+    "interactions = (\n",
+    "    -jnp.array([0, 0.001, 0.001]),\n",
+    "    -jnp.array([0.001, 0, 0.001]),\n",
+    "    -jnp.array([0.001, 0.001, 0]),\n",
+    ")"
    ]
   },
   {
@@ -286,9 +285,21 @@
    "source": [
     "vectorized_u_init = vmap(lambda t: u(t, init_params), (0), 0)\n",
     "\n",
-    "plt.plot(train_data.times.sort(axis=0) * Tmax, vectorized_u_init(train_data.times.sort(axis=0))[:, 0], label=\"N1\")\n",
-    "plt.plot(train_data.times.sort(axis=0) * Tmax, vectorized_u_init(train_data.times.sort(axis=0))[:, 1], label=\"N2\")\n",
-    "plt.plot(train_data.times.sort(axis=0) * Tmax, vectorized_u_init(train_data.times.sort(axis=0))[:, 2], label=\"N3\")\n",
+    "plt.plot(\n",
+    "    train_data.times.sort(axis=0) * Tmax,\n",
+    "    vectorized_u_init(train_data.times.sort(axis=0))[:, 0],\n",
+    "    label=\"N1\",\n",
+    ")\n",
+    "plt.plot(\n",
+    "    train_data.times.sort(axis=0) * Tmax,\n",
+    "    vectorized_u_init(train_data.times.sort(axis=0))[:, 1],\n",
+    "    label=\"N2\",\n",
+    ")\n",
+    "plt.plot(\n",
+    "    train_data.times.sort(axis=0) * Tmax,\n",
+    "    vectorized_u_init(train_data.times.sort(axis=0))[:, 2],\n",
+    "    label=\"N3\",\n",
+    ")\n",
     "\n",
     "plt.legend()"
    ]
@@ -333,7 +344,7 @@
     "    loss_weights=loss_weights,\n",
     "    dynamic_loss=dynamic_loss,\n",
     "    initial_condition=(float(tmin), jnp.array([N_0[0], N_0[1], N_0[2]])),\n",
-    "    params=init_params\n",
+    "    params=init_params,\n",
     ")"
    ]
   },
@@ -357,10 +368,7 @@
     "train_data, batch = train_data.get_batch()\n",
     "\n",
     "losses_and_grad = jax.value_and_grad(loss.evaluate, 0, has_aux=True)\n",
-    "losses, grads = losses_and_grad(\n",
-    "    init_params,\n",
-    "    batch\n",
-    ")\n",
+    "losses, grads = losses_and_grad(init_params, batch)\n",
     "l_tot, d = losses\n",
     "print(f\"total loss: {l_tot}\")\n",
     "print(f\"Individual losses: { {key: f'{val:.2f}' for key, val in d.items()} }\")"
@@ -406,13 +414,14 @@
     "    init_value=start_learning_rate,\n",
     "    transition_begin=5000,\n",
     "    transition_steps=100,\n",
-    "    decay_rate=0.99)\n",
+    "    decay_rate=0.99,\n",
+    ")\n",
     "\n",
     "tx = optax.chain(\n",
     "    optax.scale_by_adam(),  # Use the updates from adam.\n",
     "    optax.scale_by_schedule(scheduler),  # Use the learning rate from the scheduler.\n",
     "    # Scale updates by -1 since optax.apply_updates is additive and we want to descend on the loss.\n",
-    "    optax.scale(-1.0)\n",
+    "    optax.scale(-1.0),\n",
     ")"
    ]
   },
@@ -487,12 +496,8 @@
     }
    ],
    "source": [
-    "params, total_loss_list, loss_by_term_dict, data, loss, _, _ , _, _ = jinns.solve(\n",
-    "    init_params=params,\n",
-    "    data=train_data,\n",
-    "    optimizer=tx,\n",
-    "    loss=loss,\n",
-    "    n_iter=n_iter\n",
+    "params, total_loss_list, loss_by_term_dict, data, loss, _, _, _, _, _ = jinns.solve(\n",
+    "    init_params=params, data=train_data, optimizer=tx, loss=loss, n_iter=n_iter\n",
     ")"
    ]
   },
@@ -575,20 +580,23 @@
     }
    ],
    "source": [
-    "u_est = vmap(lambda t:u(t, params), (0), 0)\n",
+    "u_est = vmap(lambda t: u(t, params), (0), 0)\n",
     "\n",
     "key, subkey = random.split(key, 2)\n",
-    "val_data = jinns.data.DataGeneratorODE(key=subkey, nt=nt, tmin=tmin, tmax=tmax, method=method)\n",
+    "val_data = jinns.data.DataGeneratorODE(\n",
+    "    key=subkey, nt=nt, tmin=tmin, tmax=tmax, method=method\n",
+    ")\n",
     "\n",
     "import pandas as pd\n",
+    "\n",
     "ts = val_data.times.sort(axis=0).squeeze()\n",
     "df = pd.DataFrame(\n",
     "    {\n",
-    "        \"t\": ts * Tmax, # rescale time for plotting\n",
+    "        \"t\": ts * Tmax,  # rescale time for plotting\n",
     "        \"N1\": u_est(ts)[:, 0],\n",
     "        \"N2\": u_est(ts)[:, 1],\n",
     "        \"N3\": u_est(ts)[:, 2],\n",
-    "        \"Method\": \"PINN\"\n",
+    "        \"Method\": \"PINN\",\n",
     "    },\n",
     ")\n",
     "df.plot(x=\"t\")"
@@ -615,6 +623,7 @@
     "\n",
     "# NOTE the following line is not accurate as it skips one batch\n",
     "\n",
+    "\n",
     "def lotka_volterra_log(y_log, t, eq_params):\n",
     "    \"\"\"\n",
     "    Generalized Lotka-Volterra model for N bacterial species, with logarithmic transformation for stability.\n",
@@ -633,14 +642,19 @@
     "    dydt = np.zeros(N)\n",
     "\n",
     "    for i in range(N):\n",
-    "        dydt[i] = y[i] * (alpha[i] - beta[i] * np.sum(y) - np.sum([gamma[j][i] * y[j] for j in range(N)]))\n",
+    "        dydt[i] = y[i] * (\n",
+    "            alpha[i]\n",
+    "            - beta[i] * np.sum(y)\n",
+    "            - np.sum([gamma[j][i] * y[j] for j in range(N)])\n",
+    "        )\n",
     "\n",
     "    dydt_log = dydt / y\n",
     "\n",
     "    return dydt_log\n",
     "\n",
+    "\n",
     "# Define name bacteria\n",
-    "names = ['N1', 'N2', 'N3']\n",
+    "names = [\"N1\", \"N2\", \"N3\"]\n",
     "N = len(names)\n",
     "\n",
     "# Define model parameters\n",
@@ -671,14 +685,11 @@
     "# comparative plots\n",
     "df_scipy = pd.DataFrame(\n",
     "    {\n",
-    "        \"t\": ts * Tmax, # rescale time for plotting\n",
-    "        \"Method\": \"Scipy solver\"\n",
-    "    } |\n",
-    "    {\n",
-    "        f\"N{i+1}\": y[:,i] for i in range(3)\n",
-    "    },\n",
-    ")\n",
-    "\n"
+    "        \"t\": ts * Tmax,  # rescale time for plotting\n",
+    "        \"Method\": \"Scipy solver\",\n",
+    "    }\n",
+    "    | {f\"N{i + 1}\": y[:, i] for i in range(3)},\n",
+    ")"
    ]
   },
   {
@@ -710,8 +721,20 @@
    ],
    "source": [
     "import seaborn as sns\n",
-    "df_plot = pd.concat((df, df_scipy)).melt(id_vars=['Method', \"t\"], var_name=\"Population\", value_name=\"Solution\")\n",
-    "sns.relplot(df_plot, kind='line', x='t', y='Solution', hue='Population',  style='Method', height=4, aspect=2)"
+    "\n",
+    "df_plot = pd.concat((df, df_scipy)).melt(\n",
+    "    id_vars=[\"Method\", \"t\"], var_name=\"Population\", value_name=\"Solution\"\n",
+    ")\n",
+    "sns.relplot(\n",
+    "    df_plot,\n",
+    "    kind=\"line\",\n",
+    "    x=\"t\",\n",
+    "    y=\"Solution\",\n",
+    "    hue=\"Population\",\n",
+    "    style=\"Method\",\n",
+    "    height=4,\n",
+    "    aspect=2,\n",
+    ")"
    ]
   },
   {
@@ -749,7 +772,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.2"
+   "version": "3.11.11"
   },
   "vscode": {
    "interpreter": {