def arm_workspace(self, arm: device.HapticDevice):
points = []
t1 = 5 * np.pi / 6
for t0 in np.linspace(np.pi / 2, -np.pi / 2, num=1000):
x, y = arm.fwd_kinematics(t0, t0 + t1)
c = calculate.Coord(cpos=(x, y),
win_width=WIN_WIDTH,
win_height=WIN_HEIGHT)
points.extend((int(c.j), int(c.i)))
t0 = -np.pi / 2
for t1 in np.linspace(5 * np.pi / 6, 0, num=1000):
x, y = arm.fwd_kinematics(t0, t0 + t1)
c = calculate.Coord(cpos=(x, y),
win_width=WIN_WIDTH,
win_height=WIN_HEIGHT)
points.extend((int(c.j), int(c.i)))
t1 = 0
for t0 in np.linspace(-np.pi / 2, np.pi / 2, num=1000):
x, y = arm.fwd_kinematics(t0, t0 + t1)
c = calculate.Coord(cpos=(x, y),
win_width=WIN_WIDTH,
win_height=WIN_HEIGHT)
points.extend((int(c.j), int(c.i)))
self.workspace_points = points
def generate_device(self, arm=device.HapticDevice(init_with_device=False)):
"""Generate device takes the haptic device class and turns it into two
arm segments that will represent the arm"""
arm0_sprite = self.spritefactory.from_color(
BLUE, size=(30, int(arm.arm0.length * calculate.PIXELS_PER_METER)))
arm0 = ArmSegment(self.world,
arm0_sprite,
wposj=WIN_WIDTH // 2,
wposi=0,
angle=0)
endr, endtheta = arm0.get_end()
pos = calculate.Coord(ppos=(endr, endtheta),
win_width=WIN_WIDTH,
win_height=WIN_HEIGHT)
arm1_sprite = self.spritefactory.from_color(
ORANGE,
size=(30, int(arm.arm1.length * calculate.PIXELS_PER_METER)))
arm1 = ArmSegment(
self.world,
arm1_sprite,
wposj=pos.window.j,
wposi=pos.window.i,
angle=0,
)
self.arms = (arm0, arm1)
def __init__(self, world, sprite, wposi=0, wposj=0, angle=0):
pos = calculate.Coord(wpos=(wposi, wposj),
win_width=WIN_WIDTH,
win_height=WIN_HEIGHT)
self.world = world
self.sprite = sprite
self.sprite.pos = pos
self.sprite.angle = angle
self.sprite.position = self.get_origin()
self.sprite.point = sdl2.SDL_Point(int(self.sprite.size[0] / 2), 0)
def vector_stream_plot(self, vf):
"""Visualizes the given vector field with a stream plot. Small
rectangles mark the sampling points and the trails move in the direction
of the vector field at those sampled points. Green rect shows the center
of the vector field."""
center = calculate.Coord(
cpos=(vf.args['xcenter'], vf.args['ycenter']),
win_width=WIN_WIDTH,
win_height=WIN_HEIGHT,
)
self.vectors = [] # stores points of vector lines
self.squares = [] # stores sampled points
self.center = [center.window.j - 5, center.window.i - 5, 10, 10]
# step size defines how much space between each sample point in number
# of pixels
step_size = 40
# iterate over the width and height of the window, includes bounds if
# step_size is a multiple of height and width
for y in np.arange(0, WIN_HEIGHT // step_size * step_size + 1,
step_size):
for x in np.arange(0, WIN_WIDTH // step_size * step_size + 1,
step_size):
# add initial point to squares and start a vector line
vector = []
coord = calculate.Coord(wpos=(y, x),
win_width=WIN_WIDTH,
win_height=WIN_HEIGHT)
vector.extend((int(coord.window.j), int(coord.window.i)))
self.squares.append((x - 3, y - 3, 6, 6))
# sample n points in direction of vf from origin and add to the
# list of vectors
for _ in range(20):
dx, dy = vf.return_vectors(*coord.cartesian)
dx /= 50
dy /= 50
coord.cartesian = (
coord.cartesian.x + dx,
coord.cartesian.y + dy,
)
vector.extend((int(coord.window.j), int(coord.window.i)))
self.vectors.append(vector)
def get_end(self):
"""gets the middle of the end of segment in polar cords"""
ox, oy = self.sprite.pos.cartesian
x = (self.sprite.size[1] / calculate.PIXELS_PER_METER) * np.cos(
np.radians(-self.sprite.angle)) + ox
y = (self.sprite.size[1] / calculate.PIXELS_PER_METER) * np.sin(
np.radians(-self.sprite.angle)) + oy
end = calculate.Coord(cpos=(x, y),
win_width=WIN_WIDTH,
win_height=WIN_HEIGHT)
return end.polar.r, end.polar.theta
def condition_number_heatmap(arm):
"""
Generates a heat map of the condition number
"""
binsz = 1
num_samples = WIN_WIDTH * WIN_HEIGHT // (binsz**2)
imat, jmat = np.meshgrid(np.arange(0, WIN_HEIGHT, binsz),
np.arange(0, WIN_WIDTH, binsz))
imat = imat.flatten()
jmat = jmat.flatten()
xmat, ymat = calculate.Coord(wpos=(imat, jmat),
win_width=WIN_WIDTH,
win_height=WIN_HEIGHT).cartesian
if not os.path.isdir('condition-number-heatmap'):
os.mkdir('condition-number-heatmap')
if os.path.isfile(f'condition-number-heatmap/data.npz'):
data = np.load(f'condition-number-heatmap/data.npz')
kmat = data['kmat']
else:
th0mat = np.empty(num_samples)
th1mat = np.empty(num_samples)
for idx, (x, y) in enumerate(zip(xmat, ymat)):
theta0, theta1 = arm.inv_kinematics_num(x, y)
th0mat[idx] = theta0
th1mat[idx] = theta1
print(f'Calculating IK: {idx+1}/{num_samples}', end='\r')
print()
kmat = np.empty(num_samples)
for idx, (th0, th1) in enumerate(zip(th0mat, th1mat)):
k = arm.condition_number(th0, th1)
kmat[idx] = k
print(
f'Calculating condition numbers: {idx+1}/{num_samples}',
end='\r',
)
print()
print('Saving data...', end='\r')
np.savez_compressed(
f'condition-number-heatmap/data',
kmat=kmat,
)
print('Saving data...done')
print('Generating condition number heatmap...', end='\r')
plt.figure(figsize=(WIN_WIDTH / 100, WIN_HEIGHT / 100))
sb.heatmap(kmat.reshape((WIN_WIDTH // binsz, WIN_HEIGHT // binsz)), )
plt.savefig(f'condition-number-heatmap/condition_number.png')
print('Generating condition number heatmap...done')
def update_device(self, theta0, theta1):
"""Updates the device configuration with the given angles"""
pos = self.arms[0].sprite.pos
angle0 = np.degrees(-theta0) # pi - theta0
self.arms[0].update(pos.window.j, pos.window.i, angle0)
r, theta = self.arms[0].get_end()
pos = calculate.Coord(ppos=(r, theta),
win_width=WIN_WIDTH,
win_height=WIN_HEIGHT)
angle1 = np.degrees(-theta1)
self.arms[1].update(pos.window.j, pos.window.i, angle1)
def step(self) -> float:
"""Advances the visualization one time/sim step ahead"""
# time frame start
frame_start = time.monotonic()
# clear old graphics and update state for next frame
self.renderer.clear()
if self.vectors is not None:
for vector in self.vectors:
self.renderer.draw_line(vector, color=WHITE)
self.renderer.draw_rect(self.squares, color=WHITE)
self.renderer.fill(self.center, color=GREEN)
# render traces
if self.trace:
p = calculate.Coord(
ppos=self.arms[1].get_end(),
win_width=WIN_WIDTH,
win_height=WIN_HEIGHT,
)
self.trace_buffer.append((int(p.j - 2), int(p.i - 2), 4, 4))
self.renderer.fill(self.trace_buffer, color=RED)
if self.workspace_points:
self.renderer.draw_line(self.workspace_points, color=CYAN)
# render arm sprites
for arm in self.arms:
arm.draw(self.spriterenderer)
# render text sprites onto window
if self.text_sprites:
self.spriterenderer.render(self.text_sprites)
self.renderer.present()
# preserve physics by calculating elapsed time and delaying if necessary
# to limit to only 60 FPS. This prevents the robot velocity being
# dependent on the frame rate of the simulation
elapsed_time = time.monotonic() - frame_start
if elapsed_time < 1 / 60:
sdl2.SDL_Delay(int((1 / 60 * 1000) - (elapsed_time * 1000)))
elapsed_time = time.monotonic() - frame_start
return elapsed_time
def main(args):
# init visualization stuff
vis = SDLWrapper(tracelen=args.trace)
vis.generate_device()
arm = device.HapticDevice(False)
# take cli args for vf
field = args.field
vf_args = {
'xcenter': args.xcenter,
'ycenter': args.ycenter,
'dtheta': args.dtheta,
'radius': args.radius,
'buffer': args.buffer,
'drmax': args.drmax,
}
vf = calculate.VectorField(arm, field=field, args=vf_args)
vis.vector_stream_plot(vf)
if args.theta_heatmap:
vis.theta_heatmap(arm, vf)
if args.condition_number_heatmap:
vis.condition_number_heatmap(arm)
if args.workspace:
vis.arm_workspace(arm)
theta0 = 0
theta1 = np.pi / 2
vis.update_device(theta0, theta1)
x, y = 0, 0
i, j = 0, 0
recalculate = False
# elapsed time variable to keep track of visualization frame rate and
# simulation rate
elapsed_time = 0
global STATE_RUNNING
STATE_RUNNING = True
while STATE_RUNNING:
if args.state == 'animate':
vis.text([['State: animate']])
t = time.monotonic()
theta0 = np.pi / 4 * np.sin(t)
theta1 = np.pi / 4 * np.cos(2 * t)
vis.update_device(theta0, theta1)
elif args.state == 'follow':
vis.text([['State: follow']])
if recalculate:
print(f'xy: ({x}, {y}) | ij: ({i}, {j})')
vis.smooth_move_to_location(arm, x, y)
recalculate = False
if len(vis.config_buffer) > 0:
vis.update_device(*vis.config_buffer.popleft())
elif len(vis.config_buffer) > 0:
vis.update_device(*vis.config_buffer.popleft())
elif args.state == 'simulate':
if recalculate:
recalculate = False
vis.smooth_move_to_location(arm, x, y)
elif len(vis.config_buffer) > 0:
theta0, theta1 = vis.config_buffer.popleft()
vis.update_device(theta0, theta1)
else:
# get x, y of end effector
x, y = arm.fwd_kinematics(theta0, theta1)
# get desired vector from vectorfield
dx, dy = vf.return_vectors(x, y)
vectors = np.array([dx, dy])
# calculate the needed angular velocities
ijac = arm.inv_jacobian(theta0, theta1)
thetas = np.matmul(ijac, vectors)
dtheta0, dtheta1 = thetas
# update arm angles
theta0 += dtheta0 * elapsed_time
theta1 += dtheta1 * elapsed_time
# update visualization
vis.update_device(theta0, theta1)
# write diagnostic text to window
text = [
['State: simulate'],
[
f'x: {x:6.3f}, y: {y:6.3f}',
],
[
f'dx: {dx:6.3f}, dy: {dy:6.3f}',
],
[
f'theta0: {theta0:6.3f}, theta1: {theta1:6.3f}',
],
[
f'dtheta0: {dtheta0:6.3f}, dtheta1: {dtheta1:6.3f}',
],
]
vis.text(text)
else:
vis.text([['State: none']])
theta0, theta1 = -np.pi / 2, -np.pi / 2 + 5 * np.pi / 6
vis.update_device(theta0, theta1)
elapsed_time = vis.step()
# check for window being closed or keypresses and keyboard control
events = sdl2.ext.get_events()
for event in events:
if event.type == sdl2.SDL_QUIT:
sdl2.ext.quit()
STATE_RUNNING = False
break
if event.type == sdl2.SDL_MOUSEMOTION:
j, i = event.motion.x, event.motion.y
if event.type == sdl2.SDL_MOUSEBUTTONDOWN:
if event.button.button == sdl2.SDL_BUTTON_LEFT:
x, y = calculate.Coord(wpos=(i, j),
win_width=WIN_WIDTH,
win_height=WIN_HEIGHT).cartesian
recalculate = True
def theta_heatmap(arm, vf):
"""
Create and update a texture that is a single color channel heatmap
representing the theta velocities.
"""
# TODO heatmap for theta velocities
# [x] calculate ik
# [x] write method to impl newton-raphson (numerical ik)
# [x] click on point in space and it calcs config and draws arm
# [ ] for range of theta0 and theta1 calc jacobian and ratio of
# eigenvalues of jacobian, smaller eigenvalue / bigger eigenvalue
# print('analytical:', arm.inv_kinematics(0, 0.4))
# print('numerical:', arm.inv_kinematics_num(0, 0.4))
binsz = 1
num_samples = WIN_WIDTH * WIN_HEIGHT // (binsz**2)
imat, jmat = np.meshgrid(np.arange(0, WIN_HEIGHT, binsz),
np.arange(0, WIN_WIDTH, binsz))
imat = imat.flatten()
jmat = jmat.flatten()
xmat, ymat = calculate.Coord(wpos=(imat, jmat),
win_width=WIN_WIDTH,
win_height=WIN_HEIGHT).cartesian
dxmat, dymat = vf.return_vectors(xmat, ymat)
if not os.path.isdir('theta-heatmap'):
os.mkdir('theta-heatmap')
if os.path.isfile(f'theta-heatmap/{vf.field}_data.npz'):
data = np.load(f'theta-heatmap/{vf.field}_data.npz')
dthetamatrix0 = data['dthetamatrix0']
dthetamatrix1 = data['dthetamatrix1']
th0mat = data['Th0']
th1mat = data['Th1']
else:
th0mat = np.empty(num_samples)
th1mat = np.empty(num_samples)
for idx, (x, y) in enumerate(zip(xmat, ymat)):
theta0, theta1 = arm.inv_kinematics_num(x, y)
th0mat[idx] = theta0
th1mat[idx] = theta1
print(f'Calculating IK: {idx+1}/{num_samples}', end='\r')
print()
dthetamatrix0 = np.empty(num_samples)
dthetamatrix1 = np.empty(num_samples)
for idx, (t0, t1, dx,
dy) in enumerate(zip(th0mat, th1mat, dxmat, dymat)):
dthetas = arm.inv_jacobian(t0, t1) @ np.array([dx, dy])
dthetamatrix0[idx] = dthetas[0]
dthetamatrix1[idx] = dthetas[1]
print(f'Calculating ijac: {idx+1}/{num_samples}', end='\r')
print()
print('Saving data...', end='\r')
np.savez_compressed(
f'theta-heatmap/{vf.field}_data',
Th0=th0mat,
Th1=th1mat,
dthetamatrix0=dthetamatrix0,
dthetamatrix1=dthetamatrix1,
)
print('Saving data...done')
print('Generating dtheta0 heatmap...', end='\r')
plt.figure(figsize=(WIN_WIDTH / 100, WIN_HEIGHT / 100))
sb.heatmap(
dthetamatrix0.reshape((WIN_WIDTH // binsz, WIN_HEIGHT // binsz)),
vmin=-2,
vmax=2,
)
# cbar=False, xticklabels=False, yticklabels=False)
plt.savefig(f'theta-heatmap/{vf.field}_dtheta0.png')
print('Generating dtheta0 heatmap...done')
print('Generating theta0 heatmap...', end='\r')
plt.figure(figsize=(WIN_WIDTH / 100, WIN_HEIGHT / 100))
sb.heatmap(
th0mat.reshape((WIN_WIDTH // binsz, WIN_HEIGHT // binsz)),
vmin=-2 * np.pi,
vmax=2 * np.pi,
)
# cbar=False, xticklabels=False, yticklabels=False)
plt.savefig(f'theta-heatmap/{vf.field}_theta0.png')
print('Generating theta0 heatmap...done')
print('Generating dtheta1 heatmap...', end='\r')
plt.figure(figsize=(WIN_WIDTH / 100, WIN_HEIGHT / 100))
sb.heatmap(
dthetamatrix1.reshape((WIN_WIDTH // binsz, WIN_HEIGHT // binsz)),
vmin=-2,
vmax=2,
)
# cbar=False, xticklabels=False, yticklabels=False)
plt.savefig(f'theta-heatmap/{vf.field}_dtheta1.png')
print('Generating dtheta1 heatmap...done')
print('Generating theta1 heatmap...', end='\r')
plt.figure(figsize=(WIN_WIDTH / 100, WIN_HEIGHT / 100))
sb.heatmap(
th1mat.reshape((WIN_WIDTH // binsz, WIN_HEIGHT // binsz)),
vmin=-2 * np.pi,
vmax=2 * np.pi,
)
# cbar=False, xticklabels=False, yticklabels=False)
plt.savefig(f'theta-heatmap/{vf.field}_theta1.png')
print('Generating theta1 heatmap...done')