fix runner by adding init_fn logic

orionarcher · orionarcher · commit 67117a884aaa · 2025-03-04T22:26:22.000Z
diff --git a/torchsim/runners.py b/torchsim/runners.py
@@ -77,13 +77,13 @@ def integrate(
         )
 
     dtype, device = state.positions.dtype, state.positions.device
-    state, update_fn = integrator(
-        state=state,
+    init_fn, update_fn = integrator(
         model=model,
         kT=torch.tensor(temps[0] * unit_system.temperature, dtype=dtype, device=device),
         dt=torch.tensor(timestep * unit_system.time, dtype=dtype, device=device),
         **integrator_kwargs,
     )
+    state = init_fn(state)
 
     for step in range(1, n_steps + 1):
         state = update_fn(state, kT=temps[step - 1] * unit_system.temperature)

Original file line number	Diff line number	Diff line change
`@@ -77,13 +77,13 @@ def integrate(`
`77`	`77`	`)`
`78`	`78`
`79`	`79`	`dtype, device = state.positions.dtype, state.positions.device`
`80`		`- state, update_fn = integrator(`
`81`		`- state=state,`
	`80`	`+ init_fn, update_fn = integrator(`
`82`	`81`	`model=model,`
`83`	`82`	`kT=torch.tensor(temps[0] * unit_system.temperature, dtype=dtype, device=device),`
`84`	`83`	`dt=torch.tensor(timestep * unit_system.time, dtype=dtype, device=device),`
`85`	`84`	`**integrator_kwargs,`
`86`	`85`	`)`
	`86`	`+ state = init_fn(state)`
`87`	`87`
`88`	`88`	`for step in range(1, n_steps + 1):`
`89`	`89`	`state = update_fn(state, kT=temps[step - 1] * unit_system.temperature)`