Small change to example

mstoelzle · mstoelzle · commit 699d262a5e5f · 2024-11-05T14:28:59.000-05:00
diff --git a/examples/simulate_planar_pcs.py b/examples/simulate_planar_pcs.py
@@ -101,6 +101,8 @@ def draw_robot(
     strain_basis, forward_kinematics_fn, dynamical_matrices_fn, auxiliary_fns = (
         planar_pcs.factory(sym_exp_filepath, strain_selector)
     )
+    # jit the functions
+    dynamical_matrices_fn = jax.jit(partial(dynamical_matrices_fn, include_coriolis=True))
     batched_forward_kinematics = vmap(
         forward_kinematics_fn, in_axes=(None, None, 0), out_axes=-1
     )
@@ -241,3 +243,4 @@ def draw_robot(
         video.write(img)
 
     video.release()
+    print(f"Video saved at {video_path}")

Original file line number	Diff line number	Diff line change
`@@ -101,6 +101,8 @@ def draw_robot(`
`101`	`101`	`strain_basis, forward_kinematics_fn, dynamical_matrices_fn, auxiliary_fns = (`
`102`	`102`	`planar_pcs.factory(sym_exp_filepath, strain_selector)`
`103`	`103`	`)`
	`104`	`+ # jit the functions`
	`105`	`+ dynamical_matrices_fn = jax.jit(partial(dynamical_matrices_fn, include_coriolis=True))`
`104`	`106`	`batched_forward_kinematics = vmap(`
`105`	`107`	`forward_kinematics_fn, in_axes=(None, None, 0), out_axes=-1`
`106`	`108`	`)`
`@@ -241,3 +243,4 @@ def draw_robot(`
`241`	`243`	`video.write(img)`
`242`	`244`
`243`	`245`	`video.release()`
	`246`	`+ print(f"Video saved at {video_path}")`