import numpy as np
import copp_py as copp
def main() -> None:
try:
import jax
import jax.numpy as jnp
except ImportError as exc:
raise SystemExit(
'Install JAX to run this example: python -m pip install "copp-py[jax]"'
) from exc
jax.config.update("jax_enable_x64", True)
dim = 3
n = 1001
def q_fn(s):
freq = jnp.array([2.0 * jnp.pi, 3.0 * jnp.pi, 5.0 * jnp.pi], dtype=jnp.float64)
phase = jnp.array([0.0, 0.3, 0.7], dtype=jnp.float64)
return jnp.sin(freq * s + phase)
path = copp.Path.from_jax(q_fn, 0.0, 1.0)
s = np.linspace(0.0, 1.0, n, dtype=np.float64)
robot = copp.Robot(dim, capacity=n)
robot.append_s(s)
robot.set_q_from_path_2nd(path, 0, n)
upper = np.ones(dim, dtype=np.float64)
lower = -upper
robot.add_velocity_limits(upper, lower, start_idx_s=0, length=n)
robot.add_acceleration_limits(upper, lower, start_idx_s=0, length=n)
problem = copp.solver.reach_set2.Problem(
robot.constraints,
idx_s_interval=(0, n - 1),
a_boundary=(0.0, 0.0),
)
options = copp.solver.reach_set2.Options()
reach_back = copp.solver.reach_set2.backward(problem, options)
reach_bidir = copp.solver.reach_set2.bidirectional(problem, options)
print("reach_set2 done.")
print(f"dim = {dim}, N = {n}")
print(
"backward-only: "
f"a_max.len() = {len(reach_back.a_max)}, "
f"a_min.len() = {len(reach_back.a_min)}"
)
print(
"bidirectional: "
f"a_max.len() = {len(reach_bidir.a_max)}, "
f"a_min.len() = {len(reach_bidir.a_min)}"
)
k0 = 0
km = n // 2
k1 = n - 1
print(
"bidirectional bounds @k=0/mid/end: "
f"[{reach_bidir.a_min[k0]:.6f}, {reach_bidir.a_max[k0]:.6f}], "
f"[{reach_bidir.a_min[km]:.6f}, {reach_bidir.a_max[km]:.6f}], "
f"[{reach_bidir.a_min[k1]:.6f}, {reach_bidir.a_max[k1]:.6f}]"
)
if __name__ == "__main__":
main()