Skip to content

Commit c88fc29

Browse files
milanofthemilanofthe
committed
Rewrite checkpoint notebook: coupled oscillators, flat style, rollback demo
1 parent 456e7ce commit c88fc29

File tree

1 file changed

+201
-28
lines changed

1 file changed

+201
-28
lines changed

‎docs/source/examples/checkpoints.ipynb‎

Lines changed: 201 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -3,12 +3,31 @@
33
{
44
"cell_type": "markdown",
55
"metadata": {},
6-
"source": "# Checkpoints\n\nPathSim supports saving and loading simulation state via checkpoints. This allows you to pause a simulation, save its complete state to disk, and resume it later from exactly where you left off. \n\nCheckpoints also enable **rollback**, where you return to a saved state and explore different what-if scenarios by changing parameters.\n\nCheckpoints use a split format: a JSON file for metadata and structure, and an NPZ file for numerical data (block states, solver histories, etc.)."
6+
"source": [
7+
"# Checkpoints\n",
8+
"\n",
9+
"PathSim supports saving and loading simulation state via checkpoints. This allows you to pause a simulation, save its complete state to disk, and resume it later from exactly where you left off.\n",
10+
"\n",
11+
"Checkpoints also enable **rollback** — returning to a saved state and exploring different what-if scenarios by changing parameters.\n",
12+
"\n",
13+
"Checkpoints use a split format: a JSON file for metadata and structure, and an NPZ file for numerical data (block states, solver histories, etc.)."
14+
]
715
},
816
{
917
"cell_type": "markdown",
1018
"metadata": {},
11-
"source": "## Setup\n\nWe'll simulate a driven harmonic oscillator — a mass-spring system excited by an external sinusoidal force. The system produces a sustained periodic response, making it easy to visually verify that checkpoints preserve continuity."
19+
"source": [
20+
"## Building the System\n",
21+
"\n",
22+
"We'll use the coupled oscillators system to demonstrate checkpoints. The energy exchange between the two oscillators produces a sustained, non-trivial response that makes it easy to visually verify checkpoint continuity."
23+
]
24+
},
25+
{
26+
"cell_type": "raw",
27+
"metadata": {},
28+
"source": [
29+
"First let's import the :class:`.Simulation` and :class:`.Connection` classes and the required blocks:"
30+
]
1231
},
1332
{
1433
"cell_type": "code",
@@ -20,76 +39,222 @@
2039
"import matplotlib.pyplot as plt\n",
2140
"\n",
2241
"from pathsim import Simulation, Connection\n",
23-
"from pathsim.blocks import Integrator, Amplifier, Adder, Scope"
42+
"from pathsim.blocks import ODE, Function, Scope"
43+
]
44+
},
45+
{
46+
"cell_type": "markdown",
47+
"metadata": {},
48+
"source": [
49+
"Define the system parameters:"
2450
]
2551
},
2652
{
2753
"cell_type": "code",
2854
"execution_count": null,
2955
"metadata": {},
3056
"outputs": [],
31-
"source": "import numpy as np\n\n# System parameters\nm = 1.0 # mass\nc = 0.1 # light damping\nk = 4.0 # spring stiffness\nF0 = 1.0 # forcing amplitude\nw = 1.8 # forcing frequency (near resonance for k/m=4 -> w0=2)\n\ndef make_system(damping=c, stiffness=k):\n \"\"\"Create a driven harmonic oscillator with configurable parameters.\"\"\"\n from pathsim.blocks import Source, Integrator, Amplifier, Adder, Scope\n\n Src = Source(lambda t: F0/m * np.sin(w * t)) # external acceleration\n I1 = Integrator(0.0) # velocity\n I2 = Integrator(0.5) # position (start displaced)\n Ac = Amplifier(-damping/m)\n Ak = Amplifier(-stiffness/m)\n P1 = Adder()\n Sc = Scope(labels=[\"position\"])\n\n blocks = [Src, I1, I2, Ac, Ak, P1, Sc]\n connections = [\n Connection(I1, I2, Ac), # velocity -> position integrator, damper\n Connection(I2, Ak, Sc), # position -> spring, scope\n Connection(Ac, P1), # -c/m * v -> adder\n Connection(Ak, P1[1]), # -k/m * x -> adder\n Connection(Src, P1[2]), # F/m -> adder\n Connection(P1, I1), # acceleration -> velocity integrator\n ]\n\n sim = Simulation(blocks, connections, dt=0.01)\n return sim, Sc"
57+
"source": [
58+
"# Mass parameters\n",
59+
"m1 = 1.0\n",
60+
"m2 = 1.5\n",
61+
"\n",
62+
"# Spring constants\n",
63+
"k1 = 2.0\n",
64+
"k2 = 3.0\n",
65+
"k12 = 0.5 # coupling spring\n",
66+
"\n",
67+
"# Damping coefficients\n",
68+
"c1 = 0.02\n",
69+
"c2 = 0.03\n",
70+
"\n",
71+
"# Initial conditions [position, velocity]\n",
72+
"x1_0 = np.array([2.0, 0.0]) # oscillator 1 displaced\n",
73+
"x2_0 = np.array([0.0, 0.0]) # oscillator 2 at rest"
74+
]
3275
},
3376
{
34-
"cell_type": "markdown",
77+
"cell_type": "raw",
3578
"metadata": {},
36-
"source": "## Save Checkpoint\n\nRun the simulation for 20 seconds, then save a checkpoint. The system will be in a sustained oscillation by this point."
79+
"source": [
80+
"Define the differential equations for each oscillator using :class:`.ODE` blocks and the coupling force using a :class:`.Function` block:"
81+
]
3782
},
3883
{
3984
"cell_type": "code",
4085
"execution_count": null,
4186
"metadata": {},
4287
"outputs": [],
43-
"source": "sim, scope = make_system()\nsim.run(20)\n\n# Save checkpoint\nsim.save_checkpoint(\"checkpoint\")\nprint(f\"Saved checkpoint at t = {sim.time:.1f}s\")"
88+
"source": [
89+
"# Oscillator 1: m1*x1'' = -k1*x1 - c1*x1' - k12*(x1 - x2)\n",
90+
"def osc1_func(x1, u, t):\n",
91+
" f_e = u[0]\n",
92+
" return np.array([x1[1], (-k1*x1[0] - c1*x1[1] - f_e) / m1])\n",
93+
"\n",
94+
"# Oscillator 2: m2*x2'' = -k2*x2 - c2*x2' + k12*(x1 - x2)\n",
95+
"def osc2_func(x2, u, t):\n",
96+
" f_e = u[0]\n",
97+
" return np.array([x2[1], (-k2*x2[0] - c2*x2[1] - f_e) / m2])\n",
98+
"\n",
99+
"# Coupling force\n",
100+
"def coupling_func(x1, x2):\n",
101+
" f = k12 * (x1 - x2)\n",
102+
" return f, -f\n",
103+
"\n",
104+
"# Blocks\n",
105+
"osc1 = ODE(osc1_func, x1_0)\n",
106+
"osc2 = ODE(osc2_func, x2_0)\n",
107+
"fn = Function(coupling_func)\n",
108+
"sc = Scope(labels=[r\"$x_1(t)$ - Oscillator 1\", r\"$x_2(t)$ - Oscillator 2\"])\n",
109+
"\n",
110+
"blocks = [osc1, osc2, fn, sc]\n",
111+
"\n",
112+
"# Connections\n",
113+
"connections = [\n",
114+
" Connection(osc1[0], fn[0], sc[0]),\n",
115+
" Connection(osc2[0], fn[1], sc[1]),\n",
116+
" Connection(fn[0], osc1[0]),\n",
117+
" Connection(fn[1], osc2[0]),\n",
118+
"]"
119+
]
44120
},
45121
{
46-
"cell_type": "markdown",
122+
"cell_type": "raw",
47123
"metadata": {},
48-
"source": "## Resume from Checkpoint\n\nLoad the checkpoint into a fresh simulation and continue for another 20 seconds. The new simulation has completely different Python objects, yet the checkpoint restores the exact state by matching blocks by type and insertion order."
124+
"source": [
125+
"Create the :class:`.Simulation` and run for 60 seconds:"
126+
]
49127
},
50128
{
51129
"cell_type": "code",
52130
"execution_count": null,
53131
"metadata": {},
54132
"outputs": [],
55-
"source": "sim_resumed, scope_resumed = make_system()\nsim_resumed.load_checkpoint(\"checkpoint\")\nprint(f\"Resumed from t = {sim_resumed.time:.1f}s\")\n\nsim_resumed.run(20)"
133+
"source": [
134+
"sim = Simulation(blocks, connections, dt=0.01)\n",
135+
"\n",
136+
"sim.run(60)\n",
137+
"\n",
138+
"fig, ax = sc.plot()\n",
139+
"plt.show()"
140+
]
141+
},
142+
{
143+
"cell_type": "markdown",
144+
"metadata": {},
145+
"source": [
146+
"The two oscillators exchange energy through the coupling spring, producing a characteristic beat pattern."
147+
]
56148
},
57149
{
58150
"cell_type": "markdown",
59151
"metadata": {},
60-
"source": "## Rollback: What-If Scenarios\n\nThis is where checkpoints really shine. We reload the same checkpoint but with **different parameters** — increasing the damping significantly. Both branches start from the exact same state at t=20, but evolve differently."
152+
"source": [
153+
"## Saving a Checkpoint\n",
154+
"\n",
155+
"Now let's save the simulation state at t=60s. This creates two files: `coupled.json` (metadata) and `coupled.npz` (numerical data)."
156+
]
61157
},
62158
{
63159
"cell_type": "code",
64160
"execution_count": null,
65161
"metadata": {},
66162
"outputs": [],
67-
"source": "# Scenario A: same parameters (continuation)\nsim_a, scope_a = make_system(damping=0.1)\nsim_a.load_checkpoint(\"checkpoint\")\nsim_a.run(20)\n\n# Scenario B: increased damping (what-if)\nsim_b, scope_b = make_system(damping=1.5)\nsim_b.load_checkpoint(\"checkpoint\")\nsim_b.run(20)\n\n# Scenario C: stiffer spring (what-if)\nsim_c, scope_c = make_system(stiffness=9.0)\nsim_c.load_checkpoint(\"checkpoint\")\nsim_c.run(20)"
163+
"source": [
164+
"sim.save_checkpoint(\"coupled\")\n",
165+
"print(f\"Checkpoint saved at t = {sim.time:.1f}s\")"
166+
]
68167
},
69168
{
70169
"cell_type": "markdown",
71170
"metadata": {},
72-
"source": "## Compare Results\n\nThe plot shows the original simulation (0–20s), followed by three different futures branching from the same checkpoint."
171+
"source": [
172+
"We can inspect the JSON file to see what was saved:"
173+
]
73174
},
74175
{
75176
"cell_type": "code",
76177
"execution_count": null,
77178
"metadata": {},
78179
"outputs": [],
79-
"source": "time_orig, data_orig = scope.read()\ntime_a, data_a = scope_a.read()\ntime_b, data_b = scope_b.read()\ntime_c, data_c = scope_c.read()\n\nfig, ax = plt.subplots(figsize=(10, 4))\n\n# Original run (0-20s)\nax.plot(time_orig, data_orig[0], \"k-\", lw=1.5, label=\"original (c=0.1, k=4)\")\n\n# Three futures from checkpoint\nax.plot(time_a, data_a[0], \"C0-\", alpha=0.8, label=\"resumed (c=0.1, k=4)\")\nax.plot(time_b, data_b[0], \"C1-\", alpha=0.8, label=\"what-if: heavy damping (c=1.5)\")\nax.plot(time_c, data_c[0], \"C2-\", alpha=0.8, label=\"what-if: stiffer spring (k=9)\")\n\nax.axvline(20, color=\"gray\", ls=\":\", alpha=0.5, lw=2, label=\"checkpoint (t=20s)\")\nax.set_xlabel(\"time [s]\")\nax.set_ylabel(\"position\")\nax.set_title(\"Checkpoint Rollback: Three Futures from the Same State\")\nax.legend(loc=\"upper right\", fontsize=8)\nfig.tight_layout()\nplt.show()"
180+
"source": [
181+
"import json\n",
182+
"\n",
183+
"with open(\"coupled.json\") as f:\n",
184+
" cp = json.load(f)\n",
185+
"\n",
186+
"print(f\"PathSim version: {cp['pathsim_version']}\")\n",
187+
"print(f\"Simulation time: {cp['simulation']['time']:.1f}s\")\n",
188+
"print(f\"Solver: {cp['simulation']['solver']}\")\n",
189+
"print(f\"Blocks saved:\")\n",
190+
"for b in cp[\"blocks\"]:\n",
191+
" print(f\" {b['_key']} ({b['type']})\")"
192+
]
193+
},
194+
{
195+
"cell_type": "raw",
196+
"metadata": {},
197+
"source": [
198+
"Blocks are identified by type and insertion order (``ODE_0``, ``ODE_1``, etc.), so the checkpoint can be loaded into any simulation with the same block structure, regardless of the specific Python objects."
199+
]
80200
},
81201
{
82202
"cell_type": "markdown",
83203
"metadata": {},
84-
"source": "All three scenarios start from the exact same state at t=20s. The blue continuation matches the original trajectory perfectly, while the heavy damping scenario (orange) decays rapidly and the stiffer spring scenario (green) shifts to a higher natural frequency."
204+
"source": [
205+
"## Rollback: What-If Scenarios\n",
206+
"\n",
207+
"This is where checkpoints really shine. We'll load the same checkpoint three times with different coupling strengths to explore how the system evolves from the exact same state.\n",
208+
"\n",
209+
"Since the checkpoint restores all block states by type and insertion order, we just need to rebuild the simulation with the same block structure but different parameters."
210+
]
211+
},
212+
{
213+
"cell_type": "code",
214+
"execution_count": null,
215+
"metadata": {},
216+
"outputs": [],
217+
"source": [
218+
"def run_scenario(k12_new, duration=60):\n",
219+
" \"\"\"Load checkpoint and continue with a different coupling constant.\"\"\"\n",
220+
" def coupling_new(x1, x2):\n",
221+
" f = k12_new * (x1 - x2)\n",
222+
" return f, -f\n",
223+
"\n",
224+
" o1 = ODE(osc1_func, x1_0)\n",
225+
" o2 = ODE(osc2_func, x2_0)\n",
226+
" f = Function(coupling_new)\n",
227+
" s = Scope()\n",
228+
"\n",
229+
" sim = Simulation(\n",
230+
" [o1, o2, f, s],\n",
231+
" [Connection(o1[0], f[0], s[0]),\n",
232+
" Connection(o2[0], f[1], s[1]),\n",
233+
" Connection(f[0], o1[0]),\n",
234+
" Connection(f[1], o2[0])],\n",
235+
" dt=0.01\n",
236+
" )\n",
237+
" sim.load_checkpoint(\"coupled\")\n",
238+
" sim.run(duration)\n",
239+
" return s.read()\n",
240+
"\n",
241+
"# Original coupling (continuation)\n",
242+
"t_a, d_a = run_scenario(k12_new=0.5)\n",
243+
"\n",
244+
"# Stronger coupling\n",
245+
"t_b, d_b = run_scenario(k12_new=2.0)\n",
246+
"\n",
247+
"# Decoupled\n",
248+
"t_c, d_c = run_scenario(k12_new=0.0)"
249+
]
85250
},
86251
{
87252
"cell_type": "markdown",
88253
"metadata": {},
89254
"source": [
90-
"## Checkpoint File Contents\n",
255+
"## Comparing the Scenarios\n",
91256
"\n",
92-
"The JSON file contains human-readable metadata about the simulation state. Let's inspect it."
257+
"The plot shows the original run (0-60s) followed by three different futures branching from the checkpoint at t=60s. We show oscillator 1 for clarity."
93258
]
94259
},
95260
{
@@ -98,24 +263,32 @@
98263
"metadata": {},
99264
"outputs": [],
100265
"source": [
101-
"import json\n",
266+
"time_orig, data_orig = sc.read()\n",
102267
"\n",
103-
"with open(\"checkpoint.json\") as f:\n",
104-
" cp = json.load(f)\n",
268+
"fig, ax = plt.subplots(figsize=(10, 4))\n",
105269
"\n",
106-
"print(f\"PathSim version: {cp['pathsim_version']}\")\n",
107-
"print(f\"Simulation time: {cp['simulation']['time']:.1f}s\")\n",
108-
"print(f\"Solver: {cp['simulation']['solver']}\")\n",
109-
"print(f\"Blocks saved: {len(cp['blocks'])}\")\n",
110-
"for b in cp[\"blocks\"]:\n",
111-
" print(f\" {b['_key']} ({b['type']})\")"
270+
"# Original run\n",
271+
"ax.plot(time_orig, data_orig[0], \"k-\", lw=1.5, label=r\"original ($k_{12}=0.5$)\")\n",
272+
"\n",
273+
"# Three futures from checkpoint\n",
274+
"ax.plot(t_a, d_a[0], \"C0-\", alpha=0.8, label=r\"continued ($k_{12}=0.5$)\")\n",
275+
"ax.plot(t_b, d_b[0], \"C1-\", alpha=0.8, label=r\"stronger coupling ($k_{12}=2.0$)\")\n",
276+
"ax.plot(t_c, d_c[0], \"C2-\", alpha=0.8, label=r\"decoupled ($k_{12}=0$)\")\n",
277+
"\n",
278+
"ax.axvline(60, color=\"gray\", ls=\":\", lw=2, alpha=0.5, label=\"checkpoint\")\n",
279+
"ax.set_xlabel(\"time [s]\")\n",
280+
"ax.set_ylabel(r\"$x_1(t)$\")\n",
281+
"ax.set_title(\"Checkpoint Rollback: Three Futures from the Same State\")\n",
282+
"ax.legend(loc=\"upper right\", fontsize=8)\n",
283+
"fig.tight_layout()\n",
284+
"plt.show()"
112285
]
113286
},
114287
{
115288
"cell_type": "markdown",
116289
"metadata": {},
117290
"source": [
118-
"Blocks are matched by type and insertion order (`Integrator_0`, `Integrator_1`, etc.), which means the checkpoint can be loaded into any simulation with the same block structure, regardless of the specific Python objects."
291+
"All three scenarios start from the exact same state at t=60s. The blue continuation matches the original trajectory perfectly, confirming checkpoint fidelity. The stronger coupling (orange) produces faster energy exchange, while the decoupled system (green) oscillates independently at its natural frequency."
119292
]
120293
}
121294
],
@@ -132,4 +305,4 @@
132305
},
133306
"nbformat": 4,
134307
"nbformat_minor": 4
135-
}
308+
}

0 commit comments

Comments
 (0)