def test_cpso_init_tree():
c = 3.
n = 4
sso = SSO("parent")
cpso = CPSO("foo")
cpso.reparentTo(sso)
objs = [RBSO(str(i)) for i in xrange(n)]
for obj in objs:
obj.setPos((c * (i - np.floor(n / 2.)), 0, 0))
obj.set_shape("Box")
cpso.add(objs)
sso.init_tree(tags=("shape", ))
assert cpso.node().getNumShapes() == n
for obj in cpso.components:
assert obj.node().getNumShapes() == 0
class Viewer(ShowBase, object):
""" Viewer for SSOs."""
def __init__(self):
ShowBase.__init__(self)
resize_window = ConfigVariableBool('viewer-resize-window', '#t')
if resize_window.getValue():
self.win_size = (800, 800)
# Black background
self.win.setClearColor((0.0, 0.0, 0.0, 1.0))
# Set up lights.
self.lights = NodePath("lights")
# Spotlight. Casts shadows.
slight = Spotlight("slight")
slight.setScene(self.render)
slight.setShadowCaster(True, 2**11, 2**11)
# Set shadow mask, so we can exclude objects from casting shadows
self.shadow_mask = BitMask32.bit(2)
slight.setCameraMask(self.shadow_mask)
slight.setColor((1.2, 1.2, 1.2, 1.))
slight.getLens().setFov(45)
slight.getLens().setNearFar(1, 100)
slnp = self.lights.attachNewNode(slight)
slnp.setPos((6, 8, 20))
slnp.lookAt(0, 0, 0)
self.render.setLight(slnp)
# Ambient light.
alight = AmbientLight("alight")
a = 0.75
alight.setColor((a, a, a, 1.0))
#alight.setColor((0.8, 0.8, 0.8, 1.0))
alnp = self.lights.attachNewNode(alight)
self.render.setLight(alnp)
self.lights.reparentTo(self.render)
# Set auto shading for shadows
use_shaders = ConfigVariableBool('viewer-use-shaders', '#t')
if use_shaders.getValue():
self.render.setShaderAuto()
# Set antialiasing on
self.render.setAntialias(AntialiasAttrib.MAuto)
# Camera
self.camera_rot = self.render.attachNewNode("camera_rot")
self.cameras = self.camera_rot.attachNewNode("cameras")
self.cameras.setPos(14, 32, 9.)
self.look_at = self.render.attachNewNode("look_at")
self.look_at.setPos(Point3(2, 0, 1))
self.cameras.lookAt(self.look_at)
self.camera.reparentTo(self.cameras)
# Adjust the camera's lens
lens = PerspectiveLens()
self.camLens = lens
self.camLens.setNearFar(0.01, 1000.0)
setlens = ConfigVariableBool('viewer-set-cam-lens', '#t')
if setlens:
self.cam.node().setLens(self.camLens)
#
# Initialize / set variables
self.sso = None
self.ssos = []
self.cache = None
self.scene = SSO("scene")
self.scene.reparentTo(self.render)
# Key callbacks.
self.accept("shift-control-escape", self.exit)
self.accept("escape", self.exit)
self.accept("0", self.reset_sso)
self.accept("arrow_left", self.prev)
self.accept("arrow_right", self.next)
self.accept("page_down", self.prev, [100])
self.accept("page_up", self.next, [100])
self.accept("f1", self.toggle_debug)
self.accept("o", self.physics_once, extraArgs=[1. / 10])
self.accept("i", self.physics_once, extraArgs=[1. / 10000])
# Remove existing keyboard tasks.
self.mandatory_events = ("window-event", "async_loader_0",
"render-texture-targets-changed",
"shift-control-escape")
# Task list: name: (key, args)
events = {
"physics": ("p", ),
"repel": ("t", ),
"bump": ("f", ),
"rotate": ("r", 20),
"rotate90": ("h", ),
"ss_task": ("s", ),
"ssa_task": ("w", ),
"bp": ("b", )
}
# Add events
for key, val in events.iteritems():
call = [key] + list(val[1:])
self.accept(val[0], self.toggle_task, call)
# These are the key events that we will never ignore
self.permanent_events = self.getAllAccepting()
# These are the key events that we will never ignore
self.permanent_tasks = [
task.getName() for task in self.taskMgr.getAllTasks()
]
self.start_time = -1
self.old_elapsed = 0
@property
def win_size(self):
""" Returns window size."""
props = WindowProperties(self.win.getProperties())
return props.getXSize(), props.getYSize()
@win_size.setter
def win_size(self, wh):
""" Sets window size."""
props = WindowProperties(self.win.getProperties())
props.setSize(*wh)
self.size = wh
self.win.requestProperties(props)
def toggle_fullscreen(self):
""" Toggles fullscreen mode."""
props = WindowProperties(self.win.getProperties())
if props.getFullscreen():
props.setSize(*self.size)
props.setFullscreen(False)
else:
w = self.pipe.getDisplayWidth()
h = self.pipe.getDisplayHeight()
props.setSize(w, h)
props.setFullscreen(True)
self.win.requestProperties(props)
def _get_screen_size(self):
winx = self.win.getXSize()
winy = self.win.getYSize()
return winx, winy
def _convert_coordinate(self, P0):
""" Convert 3 coordinates to 2d projection, and 2d coordinates
to 3d extrusion."""
P0 = array(P0)
proj_mat = get_projection_mat(self.cam)
if P0.size == 2:
# 2d to 3d.
line = extrude(P0, proj_mat)
normal = array((0., 0., 1.))
P = plane_intersection(line, array(self.origin), normal)
else:
# 3d to 2d.
P = project(P0, proj_mat)
return P
def _get_screen_mouse_location(self):
""" Gets mouse location in screen coordinates."""
md = self.win.getPointer(0)
s2d = array((md.getX(), md.getY()))
return s2d
def _set_screen_mouse_location(self, s2d):
""" Sets mouse location in screen coordinates."""
self.win.movePointer(0, *s2d.astype("i"))
def _get_cursor_location(self):
""" Return cursor's 2D or 3D location."""
# Mouse's screen coordinates
x = self.mouseWatcherNode.getMouseX()
y = self.mouseWatcherNode.getMouseY()
return self._convert_coordinate((x, y))
def _set_cursor_location(self, p2d):
""" Sets cursor location in window coords [-1, 1]."""
s2d = ((p2d * array(
(1, -1)) + 1.) / 2. * array(self._get_screen_size()))
self._set_screen_mouse_location(s2d)
def _set_cursor_hidden(self, b):
""" Toggle cursor."""
props = WindowProperties()
props.setCursorHidden(b)
self.win.requestProperties(props)
def draw_cursor2d(self, task):
""" Draw cursor indicator."""
if getattr(self, "cursor", None) and self.mouseWatcherNode.hasMouse():
res = self._get_screen_size()
ar = float(res[0]) / res[1]
mx = self.mouseWatcherNode.getMouseX()
my = self.mouseWatcherNode.getMouseY()
self.cursor.setPos(mx * ar, 0, my)
return task.cont
# def draw_cursor2d(self, task):
# """ Draw cursor indicator."""
# if getattr(self, "cursor", None) and self.mouseWatcherNode.hasMouse():
# mx = self.mouseWatcherNode.getMouseX()
# my = self.mouseWatcherNode.getMouseY()
# p3d = self._convert_coordinate((mx, my))
# p2d = self._convert_coordinate(p3d)[0].squeeze()
# res = self._get_screen_size()
# ar = float(res[0]) / res[1]
# x = p2d[0] * ar
# y = p2d[1]
# self.cursor.setPos(x, 0., y)
# return task.cont
def init_physics(self, bbase):
""" Initialize the physics resources."""
self.bbase = bbase
self.debug_np = self.render.attachNewNode(self.bbase.setup_debug())
def init_ssos(self, ssos):
""" Initialize the ssos."""
GSO.loader = Loader # self.graphicsEngine.getDefaultLoader()
# Put all the input ssos into one list.
self.ssos = []
for sso in ssos:
if not isinstance(sso, NodePath):
raise TypeError("Must be NodePath: %s (%s)" % (sso, type(sso)))
# Set up the node and its descendants.
sso.init_tree(tags=("model", ))
self.ssos.append(sso)
# Number of ssos.
self.n_ssos = len(self.ssos)
def init_background(self, bg):
""" Initialize the background."""
# Put all the input ssos into one list.
if not isinstance(bg, NodePath):
raise TypeError("Must be NodePath: %s (%s)" % (bg, type(bg)))
GSO.loader = Loader # self.graphicsEngine.getDefaultLoader()
bg.init_tree(tags=("model", ))
self.background = bg
self.background.reparentTo(self.scene)
def optimize_camera(self):
""" Calculate good camera parameters given the current stim."""
top = self.cameras.getTop()
p0 = Point3()
p1 = Point3()
self.sso.calcTightBounds(p0, p1)
shape = p1 - p0
extent = (shape[0], shape[2])
extent = [max(extent)] * 2
center = shape / 2. + p0
# Adjust camera's x-position.
self.cameras.setX(top, center[0])
self.cameras.setZ(top, p1[2])
# Compute where camera will point.
# look_at = Point3(center[0], self.look_at.getY(), self.look_at.getZ())
# look_at = (center[0], center[1], self.look_at.getZ())
look_at = center
origin = Point3(center[0], center[1], p1[2])
displacement = self.cameras.getPos(top) - origin
distance = displacement.length()
fov = self.cam.node().getLens().getFov()
target_ratio = 0.65
dx = extent[0] / 2. / target_ratio / tan(radians(fov[0]) / 2.)
dz = extent[1] / 2. / target_ratio / tan(radians(fov[1]) / 2.)
dr = max(dx, dz) / distance
pos = origin + displacement * dr
self.cameras.setPos(top, pos)
#BP()
# Point camera toward stim.
self.look_at.setPos(top, look_at)
self.cameras.lookAt(self.look_at)
def _load(self, model):
""" Wrapper for egg/bam loading."""
node = NodePath(GSO.loader.loadSync(model))
return node
def toggle_task(self, taskname, sort=0):
""" Toggles taskMgr task 'taskname'."""
if not self.taskMgr.hasTaskNamed(taskname):
self.taskMgr.add(getattr(self, taskname), taskname, sort=sort)
if taskname == "physics":
self.reset_physics()
else:
self.taskMgr.remove(taskname)
def reset_physics(self):
""" Resets physics."""
self.start_time = self.taskMgr.globalClock.getFrameTime()
self.old_elapsed = 0.
def physics(self, task):
""" Task: simulate physics."""
# Elapsed time.
dt = self._get_elapsed() - self.old_elapsed
# Update amount of time simulated so far.
self.old_elapsed += dt
# Step the physics dt time.
size_sub = self.bbase.sim_par["size_sub"]
n_subs = int(dt / size_sub)
self.bbase.step(dt, n_subs, size_sub)
return task.cont
def repel(self, task):
""" Task: perform repel."""
self.bbase.repel()
return task.done
def bump(self, task):
""" Task: perform bump."""
mag0 = Vec3(0, 0, 1. / self.bbase.sim_par["size"]) * 10.
pos = Point3(-1, 0, 0)
nodes = self.background.descendants()
bodies = [n.node() for n in nodes if n.type_ is BulletRigidBodyNode]
for body in bodies:
mag = mag0 * body.getMass()
print mag
body.applyForce(mag, pos)
#BP()
return task.done
def physics_once(self, dt):
""" Step the physics dt."""
n_subs = 10
size_sub = dt / n_subs
self.bbase.step(dt, n_subs, size_sub)
# self.bbase.attenuate_velocities(self.bbase.get_bodies())
def bp(self, task):
""" Task: break."""
BP()
return task.done
def toggle_debug(self):
""" Shows/hides debug node."""
if self.debug_np.isHidden():
self.debug_np.show()
else:
self.debug_np.hide()
def rotate(self, task):
""" Task: rotate camera."""
H = (self.camera_rot.getH() + 1) % 360
self.camera_rot.setH(H)
return task.cont
def rotate90(self, task):
""" Task: rotate in ticks."""
angs = [15, 105, 195, 285]
H = int(self.camera_rot.getH())
if H in angs:
self.camera_rot.setH(angs[(angs.index(H) + 1) % len(angs)])
else:
self.camera_rot.setH(angs[0])
return task.done
def ss_task(self, task):
""" Task: Take a screenshot."""
self.screenshot()
return task.done
def ssa_task(self, task):
""" Task: Take a screenshot of every sso."""
self.screenshot(namePrefix=self.sso.getName() + ".jpg",
defaultFilename=False)
if self.n_ssos - 1 == self.ssos.index(self.sso):
return task.done
self.next()
return task.cont
def _expunge_events(self):
""" Turn OFF any non-permanent key handlers."""
events = self.getAllAccepting()
for event in set(events).difference(self.permanent_events):
self.ignore(event)
def _expunge_tasks(self):
""" Turn OFF any non-permanent tasks floating around."""
tasknames = [task.getName() for task in self.taskMgr.getAllTasks()]
for taskname in set(tasknames).difference(self.permanent_tasks):
self.taskMgr.remove(taskname)
def reset_sso(self):
""" Reset to initial scene state."""
self.goto_sso(self.ssos.index(self.sso))
def goto_sso(self, i):
""" Switches to the i-th SSO."""
print "SSO %d" % i
# Remove existing tasks and events.
self._expunge_tasks()
self._expunge_events()
if getattr(self, "sso", False):
# Detach from physical world.
self.bbase.remove_all()
# Reset its state to the initial one.
self.cache.restore()
# Detach from scene.
self.sso.detachNode()
# Set the new sso.
self.sso = self.ssos[i]
self.sso.reparentTo(self.scene)
self.cache = self.scene.store_tree()
self.attach_physics()
self.optimize_camera()
def attach_physics(self):
# Attach `self.scene` to the physics world.
self.scene.init_tree(tags=("shape", ))
bnodes = self.scene.descendants(type_=PSO)
for bnode in bnodes:
bnode.setCollideMask(BitMask32.allOn())
bnode.node().setDeactivationEnabled(False)
self.bbase.attach(bnodes)
def remove_physics(self):
# Remove `self.scene` from the physics world.
self.bbase.remove(self.scene.descendants(type_=PSO))
self.scene.destroy_tree(tags=("shape", ))
def prev(self, steps=1):
""" Task: Go back one SSO."""
i = max(0, self.ssos.index(self.sso) - steps)
self.goto_sso(i)
def next(self, steps=1):
""" Task: Go forward one SSO."""
i = min(self.n_ssos - 1, self.ssos.index(self.sso) + steps)
self.goto_sso(i)
def _get_elapsed(self):
""" Gets the time spent in this phase so far."""
# Current time.
current_time = self.taskMgr.globalClock.getFrameTime()
# Elapsed time in this phase
elapsed = current_time - self.start_time
return elapsed
def run(self):
# Start with first sso.
self.goto_sso(0)
# Call parent's run().
ShowBase.run(self)
def exit(self):
""" Stuff to do before exiting."""
sys.exit()
class Simulation(Process):
"""Simulation job."""
def __init__(self, task, params, info_lock, save=False):
self.task = task
self.params = params
self.info_lock = info_lock
self.proclabel = None
self.scene = None
self.bbase = None
self.debug_np = None
self.cache = None
self.floor = None
self.cpo = None
self.posquat_sz = 7
self.save = save
self.start_time = None
self.end_time = None
self.sim_time = 0
super(Simulation, self).__init__()
@contextmanager
def _sim_context(self, pcpos):
"""Sets up the cpo."""
tags = ("shape", )
# Disconnect pcpos from tree.
parents = []
for pcpo in pcpos:
parents.append(pcpo.getParent())
NodePathCollection(pcpos).wrtReparentTo(self.scene)
self.scene.init_tree(tags=tags)
# Add the pcpos to the Bullet world.
self.bbase.attach(pcpos)
yield
# Remove the pcpos from the Bullet world.
self.bbase.remove(pcpos)
self.scene.destroy_tree(tags=tags)
# Reassemble tree.
for pcpo, parent in zip(pcpos, parents):
pcpo.wrtReparentTo(parent)
def _add_noise(self, cpos, pcpos, noises):
"""Adds geometry noise."""
if (noises == 0).all():
return
for cpo, noise in zip(cpos, noises):
pos = cpo.getPos(self.scene)
pos += Point3(*noise)
cpo.setPos(self.scene, pos)
# Repel.
with self._sim_context(pcpos):
self.bbase.repel(50)
def _prepare_resources(self):
"""Set up all of the nodes and physics resources."""
# Set up scene.
self.scene = SSO("scene")
# Physics.
self.bbase = BulletBase()
self.bbase.init()
self.bbase.gravity = self.params["physics"]["gravity"]
self.bbase.sim_par = {
"size": self.params['simulation']["step_size"],
"n_subs": self.params['simulation']['n_substeps'],
"size_sub": self.params['simulation']["substep_size"],
}
def _prepare_scene(self, cpo_path, floor_path, rec_names):
"""Sets up the current scene."""
def get_name(cpo):
return cpo.getName()
floor = load_cpo(floor_path)
PSOStyler().apply(floor, "floor")
floor.reparentTo(self.scene)
cpo = load_cpo(cpo_path)
cpo.reparentTo(self.scene)
self.cache = self.scene.store_tree()
self.floor = floor
self.cpo = cpo
self.scene.init_tree(tags=())
pcpos = self.scene.descendants(type_=PSO)
for pcpo in pcpos:
pcpo.setCollideMask(BitMask32.allOn())
pcpo.node().setDeactivationEnabled(False)
cpos_rec = self._order_cpos(cpo.descendants(type_=PSO,
names=rec_names))
return pcpos, cpos_rec
def _clean_scene(self):
"""Clean up the current scene."""
try:
self.scene.destroy_tree()
except AttributeError:
pass
try:
self.cpo.destroy_tree()
except AttributeError:
pass
self.cache = None
self.scene = None
self.cpo = None
def _clean_resources(self):
"""Clean up all resources."""
self._clean_scene()
try:
self.scene.destroy_tree()
except AttributeError:
pass
self.bbase.destroy()
def _order_cpos(self, cpos_rec0):
names0 = list(self.task["bodies"])
names1 = [cpo.getName() for cpo in cpos_rec0]
if names0 != names1:
# Names don't match.
if sorted(names0) == sorted(names1):
# If they're just out of order, fix them.
cpos_rec = [cpos_rec0[names1.index(n)] for n in names0]
else:
# Otherwise.
raise CpoMismatchError(
"Task bodies and cpos_rec mismatch:\nT: %s\nC: %s" %
(", ".join(names0), ", ".join(names1)))
else:
cpos_rec = cpos_rec0
return cpos_rec
def _set_masses(self, kappa, cpos):
blocktypes = get_blocktypes(self.cpo)
style = get_style(path(self.task["cpo_path"]).dirname().name)
styler = PSOStyler()
for cpo, blocktype in zip(cpos, blocktypes):
styler.apply(cpo, style, blocktype=blocktype, kappa=kappa)
def _simulate(self, data, noise, force, kappa, pcpos, rec_cpos, rec_ints):
# Get simulation parameters
step_size = self.bbase.sim_par["size"]
n_substeps = self.bbase.sim_par["n_subs"]
# Store pre-noise states
read(data[0], rec_cpos)
# Add position noise
self._add_noise(rec_cpos, pcpos, noise)
# Store post-noise states
read(data[1], rec_cpos)
# Update masses
self._set_masses(kappa, rec_cpos)
# Set up force function
force_dur = self.params['physics']['force_duration']
force_pcpos = [x.node() for x in rec_cpos]
force_vecpos = get_force(force['dir'], force['mag'])
condition_time = 0.
with self._sim_context(pcpos):
# Iterative over record intervals.
for i, interval in enumerate(rec_ints, start=2):
# Step size and number of substeps for this
# interval
size = interval * step_size
n_subs = interval * n_substeps
# Set force with appropriate duration.
dur = min(size, force_dur - condition_time)
if dur > 0.:
tforce = (force_pcpos, force_vecpos, dur)
else:
tforce = None
# Simulate physics for this recording interval.
self.bbase.step(size, n_subs, force=tforce)
condition_time += size
# Store the cpos' states.
read(data[i], rec_cpos)
# Sanity check, to make sure the blocks are all above
# the floor
if (data[i][..., 2] < 0).any():
mp.util.info("Object z-positions are negative!")
self.cache.restore()
return condition_time
def simulate_all(self):
## Assorted parameters.
icpo = self.task["icpo"]
conditions = self.task['conditions']
## Set up the cpo.
pcpos, record_cpos = self._prepare_scene(path(self.task["cpo_path"]),
path(self.task["floor_path"]),
self.task['bodies'])
# Determine recording intervals
record_intervals = self.task['record_intervals']
# Allocate data storage.
alldata = np.zeros(self.task['shape'])
# enumerate over the parameters of each condition
for icond, cond in enumerate(conditions):
(iS, S), (iP, P), (iK, K), (isamp, samp) = cond
data = alldata[icond]
# shape of noises is (n_sigmas, n_samples, n_objs)
noise = self.params['noises'][iS, icpo, isamp]
# shape of forces is (n_forces, n_samples, n_objs)
force = self.params['forces'][iP, icpo, isamp]
self.sim_time += self._simulate(data, noise, force, float(K),
pcpos, record_cpos,
record_intervals)
if self.save:
# Write data to file.
data_path = path(self.task["data_path"])
if not data_path.dirname().exists():
data_path.dirname().makedirs()
np.save(data_path, alldata)
# Mark simulation as complete.
self.task["complete"] = True
self._clean_scene()
def print_info(self):
self.info_lock.acquire()
n_conditions = len(self.task['conditions'])
dt = self.end_time - self.start_time
avg = timedelta(seconds=(dt.total_seconds() / float(n_conditions)))
speedup = 100 * self.sim_time / dt.total_seconds()
mp.util.info("-" * 60)
mp.util.info("Total sim. time : %s" %
str(timedelta(seconds=self.sim_time)))
mp.util.info("Total real time : %s" % str(dt))
mp.util.info("Avg. per condition: %s" % str(avg))
mp.util.info("Num conditions : %d" % n_conditions)
mp.util.info("Speedup is %.1f%%" % speedup)
mp.util.info("-" * 60)
sys.stdout.flush()
self.info_lock.release()
def run(self):
"""Run one simulation."""
self._prepare_resources()
try:
self.start_time = datetime.now()
self.simulate_all()
except KeyboardInterrupt:
mp.util.debug("Keyboard interrupt!")
sys.exit(100)
except Exception as err:
mp.util.debug("Error: %s" % err)
raise
else:
self.end_time = datetime.now()
# print out information about the task in a thread-safe
# manner (so information isn't interleaved across tasks)
self.print_info()
finally:
# Clean simulation resources.
self._clean_resources()
class Viewer(ShowBase, object):
""" Viewer for SSOs."""
def __init__(self):
ShowBase.__init__(self)
resize_window = ConfigVariableBool('viewer-resize-window', '#t')
if resize_window.getValue():
self.win_size = (800, 800)
# Black background
self.win.setClearColor((0.0, 0.0, 0.0, 1.0))
# Set up lights.
self.lights = NodePath("lights")
# Spotlight. Casts shadows.
slight = Spotlight("slight")
slight.setScene(self.render)
slight.setShadowCaster(True, 2 ** 11, 2 ** 11)
# Set shadow mask, so we can exclude objects from casting shadows
self.shadow_mask = BitMask32.bit(2)
slight.setCameraMask(self.shadow_mask)
slight.setColor((1.2, 1.2, 1.2, 1.))
slight.getLens().setFov(45)
slight.getLens().setNearFar(1, 100)
slnp = self.lights.attachNewNode(slight)
slnp.setPos((6, 8, 20))
slnp.lookAt(0, 0, 0)
self.render.setLight(slnp)
# Ambient light.
alight = AmbientLight("alight")
a = 0.75
alight.setColor((a, a, a, 1.0))
#alight.setColor((0.8, 0.8, 0.8, 1.0))
alnp = self.lights.attachNewNode(alight)
self.render.setLight(alnp)
self.lights.reparentTo(self.render)
# Set auto shading for shadows
use_shaders = ConfigVariableBool('viewer-use-shaders', '#t')
if use_shaders.getValue():
self.render.setShaderAuto()
# Set antialiasing on
self.render.setAntialias(AntialiasAttrib.MAuto)
# Camera
self.camera_rot = self.render.attachNewNode("camera_rot")
self.cameras = self.camera_rot.attachNewNode("cameras")
self.cameras.setPos(14, 32, 9.)
self.look_at = self.render.attachNewNode("look_at")
self.look_at.setPos(Point3(2, 0, 1))
self.cameras.lookAt(self.look_at)
self.camera.reparentTo(self.cameras)
# Adjust the camera's lens
lens = PerspectiveLens()
self.camLens = lens
self.camLens.setNearFar(0.01, 1000.0)
setlens = ConfigVariableBool('viewer-set-cam-lens', '#t')
if setlens:
self.cam.node().setLens(self.camLens)
#
# Initialize / set variables
self.sso = None
self.ssos = []
self.cache = None
self.scene = SSO("scene")
self.scene.reparentTo(self.render)
# Key callbacks.
self.accept("shift-control-escape", self.exit)
self.accept("escape", self.exit)
self.accept("0", self.reset_sso)
self.accept("arrow_left", self.prev)
self.accept("arrow_right", self.next)
self.accept("page_down", self.prev, [100])
self.accept("page_up", self.next, [100])
self.accept("f1", self.toggle_debug)
self.accept("o", self.physics_once, extraArgs=[1. / 10])
self.accept("i", self.physics_once, extraArgs=[1. / 10000])
# Remove existing keyboard tasks.
self.mandatory_events = ("window-event", "async_loader_0",
"render-texture-targets-changed",
"shift-control-escape")
# Task list: name: (key, args)
events = {"physics": ("p",),
"repel": ("t",),
"bump": ("f",),
"rotate": ("r", 20),
"rotate90": ("h",),
"ss_task": ("s",),
"ssa_task": ("w",),
"bp": ("b",)}
# Add events
for key, val in events.iteritems():
call = [key] + list(val[1:])
self.accept(val[0], self.toggle_task, call)
# These are the key events that we will never ignore
self.permanent_events = self.getAllAccepting()
# These are the key events that we will never ignore
self.permanent_tasks = [task.getName()
for task in self.taskMgr.getAllTasks()]
self.start_time = -1
self.old_elapsed = 0
@property
def win_size(self):
""" Returns window size."""
props = WindowProperties(self.win.getProperties())
return props.getXSize(), props.getYSize()
@win_size.setter
def win_size(self, wh):
""" Sets window size."""
props = WindowProperties(self.win.getProperties())
props.setSize(*wh)
self.size = wh
self.win.requestProperties(props)
def toggle_fullscreen(self):
""" Toggles fullscreen mode."""
props = WindowProperties(self.win.getProperties())
if props.getFullscreen():
props.setSize(*self.size)
props.setFullscreen(False)
else:
w = self.pipe.getDisplayWidth()
h = self.pipe.getDisplayHeight()
props.setSize(w, h)
props.setFullscreen(True)
self.win.requestProperties(props)
def _get_screen_size(self):
winx = self.win.getXSize()
winy = self.win.getYSize()
return winx, winy
def _convert_coordinate(self, P0):
""" Convert 3 coordinates to 2d projection, and 2d coordinates
to 3d extrusion."""
P0 = array(P0)
proj_mat = get_projection_mat(self.cam)
if P0.size == 2:
# 2d to 3d.
line = extrude(P0, proj_mat)
normal = array((0., 0., 1.))
P = plane_intersection(line, array(self.origin), normal)
else:
# 3d to 2d.
P = project(P0, proj_mat)
return P
def _get_screen_mouse_location(self):
""" Gets mouse location in screen coordinates."""
md = self.win.getPointer(0)
s2d = array((md.getX(), md.getY()))
return s2d
def _set_screen_mouse_location(self, s2d):
""" Sets mouse location in screen coordinates."""
self.win.movePointer(0, *s2d.astype("i"))
def _get_cursor_location(self):
""" Return cursor's 2D or 3D location."""
# Mouse's screen coordinates
x = self.mouseWatcherNode.getMouseX()
y = self.mouseWatcherNode.getMouseY()
return self._convert_coordinate((x, y))
def _set_cursor_location(self, p2d):
""" Sets cursor location in window coords [-1, 1]."""
s2d = ((p2d * array((1, -1)) + 1.) / 2. *
array(self._get_screen_size()))
self._set_screen_mouse_location(s2d)
def _set_cursor_hidden(self, b):
""" Toggle cursor."""
props = WindowProperties()
props.setCursorHidden(b)
self.win.requestProperties(props)
def draw_cursor2d(self, task):
""" Draw cursor indicator."""
if getattr(self, "cursor", None) and self.mouseWatcherNode.hasMouse():
res = self._get_screen_size()
ar = float(res[0]) / res[1]
mx = self.mouseWatcherNode.getMouseX()
my = self.mouseWatcherNode.getMouseY()
self.cursor.setPos(mx * ar, 0, my)
return task.cont
# def draw_cursor2d(self, task):
# """ Draw cursor indicator."""
# if getattr(self, "cursor", None) and self.mouseWatcherNode.hasMouse():
# mx = self.mouseWatcherNode.getMouseX()
# my = self.mouseWatcherNode.getMouseY()
# p3d = self._convert_coordinate((mx, my))
# p2d = self._convert_coordinate(p3d)[0].squeeze()
# res = self._get_screen_size()
# ar = float(res[0]) / res[1]
# x = p2d[0] * ar
# y = p2d[1]
# self.cursor.setPos(x, 0., y)
# return task.cont
def init_physics(self, bbase):
""" Initialize the physics resources."""
self.bbase = bbase
self.debug_np = self.render.attachNewNode(self.bbase.setup_debug())
def init_ssos(self, ssos):
""" Initialize the ssos."""
GSO.loader = Loader # self.graphicsEngine.getDefaultLoader()
# Put all the input ssos into one list.
self.ssos = []
for sso in ssos:
if not isinstance(sso, NodePath):
raise TypeError("Must be NodePath: %s (%s)" % (sso, type(sso)))
# Set up the node and its descendants.
sso.init_tree(tags=("model",))
self.ssos.append(sso)
# Number of ssos.
self.n_ssos = len(self.ssos)
def init_background(self, bg):
""" Initialize the background."""
# Put all the input ssos into one list.
if not isinstance(bg, NodePath):
raise TypeError("Must be NodePath: %s (%s)" % (bg, type(bg)))
GSO.loader = Loader # self.graphicsEngine.getDefaultLoader()
bg.init_tree(tags=("model",))
self.background = bg
self.background.reparentTo(self.scene)
def optimize_camera(self):
""" Calculate good camera parameters given the current stim."""
top = self.cameras.getTop()
p0 = Point3()
p1 = Point3()
self.sso.calcTightBounds(p0, p1)
shape = p1 - p0
extent = (shape[0], shape[2])
extent = [max(extent)] * 2
center = shape / 2. + p0
# Adjust camera's x-position.
self.cameras.setX(top, center[0])
self.cameras.setZ(top, p1[2])
# Compute where camera will point.
# look_at = Point3(center[0], self.look_at.getY(), self.look_at.getZ())
# look_at = (center[0], center[1], self.look_at.getZ())
look_at = center
origin = Point3(center[0], center[1], p1[2])
displacement = self.cameras.getPos(top) - origin
distance = displacement.length()
fov = self.cam.node().getLens().getFov()
target_ratio = 0.65
dx = extent[0] / 2. / target_ratio / tan(radians(fov[0]) / 2.)
dz = extent[1] / 2. / target_ratio / tan(radians(fov[1]) / 2.)
dr = max(dx, dz) / distance
pos = origin + displacement * dr
self.cameras.setPos(top, pos)
#BP()
# Point camera toward stim.
self.look_at.setPos(top, look_at)
self.cameras.lookAt(self.look_at)
def _load(self, model):
""" Wrapper for egg/bam loading."""
node = NodePath(GSO.loader.loadSync(model))
return node
def toggle_task(self, taskname, sort=0):
""" Toggles taskMgr task 'taskname'."""
if not self.taskMgr.hasTaskNamed(taskname):
self.taskMgr.add(getattr(self, taskname), taskname, sort=sort)
if taskname == "physics":
self.reset_physics()
else:
self.taskMgr.remove(taskname)
def reset_physics(self):
""" Resets physics."""
self.start_time = self.taskMgr.globalClock.getFrameTime()
self.old_elapsed = 0.
def physics(self, task):
""" Task: simulate physics."""
# Elapsed time.
dt = self._get_elapsed() - self.old_elapsed
# Update amount of time simulated so far.
self.old_elapsed += dt
# Step the physics dt time.
size_sub = self.bbase.sim_par["size_sub"]
n_subs = int(dt / size_sub)
self.bbase.step(dt, n_subs, size_sub)
return task.cont
def repel(self, task):
""" Task: perform repel."""
self.bbase.repel()
return task.done
def bump(self, task):
""" Task: perform bump."""
mag0 = Vec3(0, 0, 1. / self.bbase.sim_par["size"]) * 10.
pos = Point3(-1, 0, 0)
nodes = self.background.descendants()
bodies = [n.node() for n in nodes if n.type_ is BulletRigidBodyNode]
for body in bodies:
mag = mag0 * body.getMass()
print mag
body.applyForce(mag, pos)
#BP()
return task.done
def physics_once(self, dt):
""" Step the physics dt."""
n_subs = 10
size_sub = dt / n_subs
self.bbase.step(dt, n_subs, size_sub)
# self.bbase.attenuate_velocities(self.bbase.get_bodies())
def bp(self, task):
""" Task: break."""
BP()
return task.done
def toggle_debug(self):
""" Shows/hides debug node."""
if self.debug_np.isHidden():
self.debug_np.show()
else:
self.debug_np.hide()
def rotate(self, task):
""" Task: rotate camera."""
H = (self.camera_rot.getH() + 1) % 360
self.camera_rot.setH(H)
return task.cont
def rotate90(self, task):
""" Task: rotate in ticks."""
angs = [15, 105, 195, 285]
H = int(self.camera_rot.getH())
if H in angs:
self.camera_rot.setH(angs[(angs.index(H) + 1) % len(angs)])
else:
self.camera_rot.setH(angs[0])
return task.done
def ss_task(self, task):
""" Task: Take a screenshot."""
self.screenshot()
return task.done
def ssa_task(self, task):
""" Task: Take a screenshot of every sso."""
self.screenshot(namePrefix=self.sso.getName() + ".jpg",
defaultFilename=False)
if self.n_ssos - 1 == self.ssos.index(self.sso):
return task.done
self.next()
return task.cont
def _expunge_events(self):
""" Turn OFF any non-permanent key handlers."""
events = self.getAllAccepting()
for event in set(events).difference(self.permanent_events):
self.ignore(event)
def _expunge_tasks(self):
""" Turn OFF any non-permanent tasks floating around."""
tasknames = [task.getName() for task in self.taskMgr.getAllTasks()]
for taskname in set(tasknames).difference(self.permanent_tasks):
self.taskMgr.remove(taskname)
def reset_sso(self):
""" Reset to initial scene state."""
self.goto_sso(self.ssos.index(self.sso))
def goto_sso(self, i):
""" Switches to the i-th SSO."""
print "SSO %d" % i
# Remove existing tasks and events.
self._expunge_tasks()
self._expunge_events()
if getattr(self, "sso", False):
# Detach from physical world.
self.bbase.remove_all()
# Reset its state to the initial one.
self.cache.restore()
# Detach from scene.
self.sso.detachNode()
# Set the new sso.
self.sso = self.ssos[i]
self.sso.reparentTo(self.scene)
self.cache = self.scene.store_tree()
self.attach_physics()
self.optimize_camera()
def attach_physics(self):
# Attach `self.scene` to the physics world.
self.scene.init_tree(tags=("shape",))
bnodes = self.scene.descendants(type_=PSO)
for bnode in bnodes:
bnode.setCollideMask(BitMask32.allOn())
bnode.node().setDeactivationEnabled(False)
self.bbase.attach(bnodes)
def remove_physics(self):
# Remove `self.scene` from the physics world.
self.bbase.remove(self.scene.descendants(type_=PSO))
self.scene.destroy_tree(tags=("shape",))
def prev(self, steps=1):
""" Task: Go back one SSO."""
i = max(0, self.ssos.index(self.sso) - steps)
self.goto_sso(i)
def next(self, steps=1):
""" Task: Go forward one SSO."""
i = min(self.n_ssos - 1, self.ssos.index(self.sso) + steps)
self.goto_sso(i)
def _get_elapsed(self):
""" Gets the time spent in this phase so far."""
# Current time.
current_time = self.taskMgr.globalClock.getFrameTime()
# Elapsed time in this phase
elapsed = current_time - self.start_time
return elapsed
def run(self):
# Start with first sso.
self.goto_sso(0)
# Call parent's run().
ShowBase.run(self)
def exit(self):
""" Stuff to do before exiting."""
sys.exit()