def generate_alphabet_image(self):
self.goal_characters_image = {}
for c in self.goal_characters:
default_config = self.get_default_config(c=c)
goal_c_pos = self.goal_characters_position[c]
self.set_scene(default_config)
all_positions = pyflex.get_positions().reshape([-1, 4])
all_positions = goal_c_pos.copy() # ignore the first a few cloth particles
pyflex.set_positions(all_positions)
self.update_camera('default_camera', default_config['camera_params']['default_camera']) # why we need to do this?
self.action_tool.reset([0., -1., 0.]) # hide picker
# goal_c_img = self.get_image(self.camera_height, self.camera_width)
# import time
# for _ in range(50):
# pyflex.step(render=True)
# time.sleep(0.1)
# cv2.imshow('img', self.get_image())
# cv2.waitKey()
goal_c_img = self.get_image(self.camera_height, self.camera_width)
# cv2.imwrite('../data/images/rope-configuration-goal-image-{}.png'.format(c), goal_c_img[:,:,::-1])
# exit()
self.goal_characters_image[c] = goal_c_img.copy()
def random_pick_and_place(pick_num=10, pick_scale=0.01):
""" Random pick a particle up and the drop it for pick_num times"""
curr_pos = pyflex.get_positions().reshape(-1, 4)
num_particles = curr_pos.shape[0]
for i in range(pick_num):
pick_id = np.random.randint(num_particles)
pick_dir = np.random.random(3) * 2 - 1
pick_dir[1] = (pick_dir[1] + 1)
pick_dir *= pick_scale
original_inv_mass = curr_pos[pick_id, 3]
for _ in range(60):
curr_pos = pyflex.get_positions().reshape(-1, 4)
curr_pos[pick_id, :3] += pick_dir
curr_pos[pick_id, 3] = 0
pyflex.set_positions(curr_pos.flatten())
pyflex.step()
# Revert mass
curr_pos = pyflex.get_positions().reshape(-1, 4)
curr_pos[pick_id, 3] = original_inv_mass
pyflex.set_positions(curr_pos.flatten())
pyflex.step()
# Wait to stabalize
for _ in range(100):
pyflex.step()
curr_vel = pyflex.get_velocities()
if np.alltrue(curr_vel < 0.01):
break
for _ in range(500):
pyflex.step()
curr_vel = pyflex.get_velocities()
if np.alltrue(curr_vel < 0.01):
break
def set_state(self, state_dict):
pyflex.set_positions(state_dict['particle_pos'])
pyflex.set_velocities(state_dict['particle_vel'])
pyflex.set_shape_states(state_dict['shape_pos'])
pyflex.set_phases(state_dict['phase'])
self.camera_params = copy.deepcopy(state_dict['camera_params'])
self.update_camera(self.camera_name)
def _set_to_vertical(self, x_low, height_low):
curr_pos = pyflex.get_positions().reshape((-1, 4))
vertical_pos = self._get_vertical_pos(x_low, height_low)
curr_pos[:, :3] = vertical_pos
max_height = np.max(curr_pos[:, 1])
if max_height < 0.5:
curr_pos[:, 1] += 0.5 - max_height
pyflex.set_positions(curr_pos)
pyflex.step()
def center_object():
"""
Center the object to be at the origin
NOTE: call a pyflex.set_positions and then pyflex.step
"""
pos = pyflex.get_positions().reshape(-1, 4)
pos[:, [0, 2]] -= np.mean(pos[:, [0, 2]], axis=0, keepdims=True)
pyflex.set_positions(pos.flatten())
pyflex.step()
def rotate_particles(self, angle):
pos = pyflex.get_positions().reshape(-1, 4)
center = np.mean(pos, axis=0)
pos -= center
new_pos = pos.copy()
new_pos[:, 0] = (np.cos(angle) * pos[:, 0] - np.sin(angle) * pos[:, 2])
new_pos[:, 2] = (np.sin(angle) * pos[:, 0] + np.cos(angle) * pos[:, 2])
new_pos += center
pyflex.set_positions(new_pos)
def set_state(self, state_dic):
'''
set the postion, velocity of flex particles, and postions of flex shapes.
'''
self.box_params = state_dic['box_params']
pyflex.set_positions(state_dic["particle_pos"])
pyflex.set_velocities(state_dic["particle_vel"])
pyflex.set_shape_states(state_dic["shape_pos"])
self.box_x = state_dic['box_x']
self.box_states = state_dic['box_states']
for _ in range(5):
pyflex.step()
def set_state(self, state_dic):
'''
set the postion, velocity of flex particles, and postions of flex shapes.
'''
pyflex.set_positions(state_dic["particle_pos"])
pyflex.set_velocities(state_dic["particle_vel"])
pyflex.set_shape_states(state_dic["shape_pos"])
self.glass_x = state_dic['glass_x']
self.glass_y = state_dic['glass_y']
self.glass_rotation = state_dic['glass_rotation']
self.glass_states = state_dic['glass_states']
self.poured_glass_states = state_dic['poured_glass_states']
for _ in range(5):
pyflex.step()
def generate_env_variation(self, num_variations=2, vary_cloth_size=True):
""" Generate initial states. Note: This will also change the current states! """
max_wait_step = 1000 # Maximum number of steps waiting for the cloth to stablize
stable_vel_threshold = 0.2 # Cloth stable when all particles' vel are smaller than this
generated_configs, generated_states = [], []
default_config = self.get_default_config()
default_config['flip_mesh'] = 1
for i in range(num_variations):
config = deepcopy(default_config)
self.update_camera(config['camera_name'],
config['camera_params'][config['camera_name']])
if vary_cloth_size:
cloth_dimx, cloth_dimy = self._sample_cloth_size()
config['ClothSize'] = [cloth_dimx, cloth_dimy]
else:
cloth_dimx, cloth_dimy = config['ClothSize']
self.set_scene(config)
self.action_tool.reset([0., -1., 0.])
pos = pyflex.get_positions().reshape(-1, 4)
pos[:, :3] -= np.mean(pos, axis=0)[:3]
if self.action_mode in ['sawyer', 'franka'
]: # Take care of the table in robot case
pos[:, 1] = 0.57
else:
pos[:, 1] = 0.005
pos[:, 3] = 1
pyflex.set_positions(pos.flatten())
pyflex.set_velocities(np.zeros_like(pos))
for _ in range(5): # In case if the cloth starts in the air
pyflex.step()
for wait_i in range(max_wait_step):
pyflex.step()
curr_vel = pyflex.get_velocities()
if np.alltrue(np.abs(curr_vel) < stable_vel_threshold):
break
center_object()
angle = (np.random.random() - 0.5) * np.pi / 2
self.rotate_particles(angle)
generated_configs.append(deepcopy(config))
print('config {}: {}'.format(i, config['camera_params']))
generated_states.append(deepcopy(self.get_state()))
return generated_configs, generated_states
def _set_to_flatten(self):
# self._get_current_covered_area(pyflex.get_positions().reshape(-))
cloth_dimx, cloth_dimz = self.get_current_config()['ClothSize']
N = cloth_dimx * cloth_dimz
px = np.linspace(0, cloth_dimx * self.cloth_particle_radius,
cloth_dimx)
py = np.linspace(0, cloth_dimz * self.cloth_particle_radius,
cloth_dimz)
xx, yy = np.meshgrid(px, py)
new_pos = np.empty(shape=(N, 4), dtype=np.float)
new_pos[:, 0] = xx.flatten()
new_pos[:, 1] = self.cloth_particle_radius
new_pos[:, 2] = yy.flatten()
new_pos[:, 3] = 1.
new_pos[:, :3] -= np.mean(new_pos[:, :3], axis=0)
pyflex.set_positions(new_pos.flatten())
return self._get_current_covered_area(new_pos)
def generate_env_variation(self, num_variations=1, vary_cloth_size=True):
""" Generate initial states. Note: This will also change the current states! """
max_wait_step = 500 # Maximum number of steps waiting for the cloth to stablize
stable_vel_threshold = 0.1 # Cloth stable when all particles' vel are smaller than this
generated_configs, generated_states = [], []
default_config = self.get_default_config()
for i in range(num_variations):
config = deepcopy(default_config)
self.update_camera(config['camera_name'],
config['camera_params'][config['camera_name']])
if vary_cloth_size:
cloth_dimx, cloth_dimy = self._sample_cloth_size()
config['ClothSize'] = [cloth_dimx, cloth_dimy]
else:
cloth_dimx, cloth_dimy = config['ClothSize']
self.set_scene(config)
self.action_tool.reset([0., -1., 0.])
pickpoints = self._get_drop_point_idx(
)[:2] # Pick two corners of the cloth and wait until stablize
config['target_pos'] = self._get_flat_pos()
self._set_to_vertical(x_low=np.random.random() * 0.2 - 0.1,
height_low=np.random.random() * 0.1 + 0.1)
# Get height of the cloth without the gravity. With gravity, it will be longer
p1, _, p2, _ = self._get_key_point_idx()
curr_pos = pyflex.get_positions().reshape(-1, 4)
curr_pos[0] += np.random.random() * 0.001 # Add small jittering
original_inv_mass = curr_pos[pickpoints, 3]
curr_pos[
pickpoints,
3] = 0 # Set mass of the pickup point to infinity so that it generates enough force to the rest of the cloth
pickpoint_pos = curr_pos[pickpoints, :3]
pyflex.set_positions(curr_pos.flatten())
picker_radius = self.action_tool.picker_radius
self.action_tool.update_picker_boundary([-0.3, 0.05, -0.5],
[0.5, 2, 0.5])
self.action_tool.set_picker_pos(picker_pos=pickpoint_pos +
np.array([0., picker_radius, 0.]))
# Pick up the cloth and wait to stablize
for j in range(0, max_wait_step):
pyflex.step()
curr_pos = pyflex.get_positions().reshape((-1, 4))
curr_vel = pyflex.get_velocities().reshape((-1, 3))
if np.alltrue(curr_vel < stable_vel_threshold) and j > 300:
break
curr_pos[pickpoints, :3] = pickpoint_pos
pyflex.set_positions(curr_pos)
curr_pos = pyflex.get_positions().reshape((-1, 4))
curr_pos[pickpoints, 3] = original_inv_mass
pyflex.set_positions(curr_pos.flatten())
generated_configs.append(deepcopy(config))
print('config {}: {}'.format(i, config['camera_params']))
generated_states.append(deepcopy(self.get_state()))
return generated_configs, generated_states
def _set_to_folded(self):
config = self.get_current_config()
num_particles = np.prod(config['ClothSize'], dtype=int)
particle_grid_idx = np.array(list(range(num_particles))).reshape(
config['ClothSize'][1],
config['ClothSize'][0]) # Reversed index here
cloth_dimx = config['ClothSize'][0]
x_split = cloth_dimx // 2
fold_group_a = particle_grid_idx[:, :x_split].flatten()
fold_group_b = np.flip(particle_grid_idx,
axis=1)[:, :x_split].flatten()
curr_pos = pyflex.get_positions().reshape((-1, 4))
curr_pos[fold_group_a, :] = curr_pos[fold_group_b, :].copy()
curr_pos[
fold_group_a,
1] += 0.05 # group a particle position made tcurr_pos[self.fold_group_b, 1] + 0.05e at top of group b position.
pyflex.set_positions(curr_pos)
for i in range(10):
pyflex.step()
return self._get_info()['performance']
def step(self, action):
""" Action is in 5D: (u,v) the start of the pick in image coordinate; (dx, dy, dz): the relative position of the place w.r.t. the pick"""
u, v = action[:2]
u = ((u + 1.) * 0.5) * self.image_size[0]
v = ((v + 1.) * 0.5) * self.image_size[1]
x, y, z = self._get_world_coor_from_image(u, v)
y += 0.01
dx, dy, dz = action[2:]
st_high = np.array([x, 0.2, z, 0])
st = np.array([x, y, z, 0])
en = st + np.array([dx, dy, dz, 1])
# print('st:', st)
if self.full:
self.total_steps += super().step(st_high)
self.total_steps += super().step(st)
self.total_steps += super().step(en)
en[3] = 0 # Drop cloth
# Unpick all particles
_, particle_pos = self._get_pos()
new_particle_pos = particle_pos.copy()
for i in range(self.num_picker):
if self.picked_particles[i] is not None:
new_particle_pos[
self.picked_particles[i], 3] = self.particle_inv_mass[
self.picked_particles[i]] # Revert the mass
self.picked_particles[i] = None
pyflex.set_positions(new_particle_pos)
for i in range(20):
pyflex.step()
if self.env is not None and self.env.recording:
self.env.video_frames.append(
self.env.render(mode='rgb_array'))
self.total_steps += 20
else:
raise NotImplementedError
return self.total_steps
v_box += rand_float(-0.15, 0.15) - x_box * 0.1
shape_states_ = calc_shape_states(x_box, x_box_last, scene_params[-2:])
pyflex.set_shape_states(shape_states_)
positions[i] = pyflex.get_positions().reshape(-1, dim_position)
velocities[i] = pyflex.get_velocities().reshape(-1, dim_velocity)
shape_states[i] = pyflex.get_shape_states().reshape(-1, dim_shape_state)
if i == 0:
print(np.min(positions[i], 0), np.max(positions[i], 0))
print(x_box, box_dis_x, box_dis_z)
pyflex.step()
# playback
pyflex.set_scene(6, scene_params, 0)
for i in range(len(boxes) - 1):
halfEdge = boxes[i][0]
center = boxes[i][1]
quat = boxes[i][2]
pyflex.add_box(halfEdge, center, quat)
for i in range(time_step):
pyflex.set_positions(positions[i])
pyflex.set_shape_states(shape_states[i, :-1])
pyflex.render(capture=1, path=os.path.join(des_dir, 'render_%d.tga' % i))
pyflex.clean()
def _set_to_flat(self):
curr_pos = pyflex.get_positions().reshape((-1, 4))
flat_pos = self._get_flat_pos()
curr_pos[:, :3] = flat_pos
pyflex.set_positions(curr_pos)
pyflex.step()
# Set all the particles to unseen locations
def clear_pos():
pos[:, 0:3] = np.ones((num_particles, 3)) * 20
# Function for putting particle positions into the array
def load_pos(particles):
print(particles.shape[0])
for i in range(particles.shape[0]):
if i < num_particles:
pos[i, 0:3] = particles[i, 0:3]
# # Load in the particles
# path = particle_path_prefix+"frame_{}.npy".format(599)
# particles = np.load(path)
# load_pos(particles)
# pyflex.set_positions(pos)
for i in range(time_step):
# Load in the particles
path = particle_path_prefix + "frame_{}.npy".format(i)
particles = np.load(path) / scale
clear_pos()
load_pos(particles)
pyflex.set_positions(pos)
# pyflex.step(capture=1, path=os.path.join(des_dir, "render_%04d.tga" % i))
pyflex.render(capture=1, path=os.path.join(des_dir, "render_%04d.tga" % i))
pyflex.clean()
def generate_env_variation(self, num_variations=1, vary_cloth_size=True):
""" Generate initial states. Note: This will also change the current states! """
max_wait_step = 300 # Maximum number of steps waiting for the cloth to stablize
stable_vel_threshold = 0.01 # Cloth stable when all particles' vel are smaller than this
generated_configs, generated_states = [], []
default_config = self.get_default_config()
for i in range(num_variations):
config = deepcopy(default_config)
self.update_camera(config['camera_name'],
config['camera_params'][config['camera_name']])
if vary_cloth_size:
cloth_dimx, cloth_dimy = self._sample_cloth_size()
config['ClothSize'] = [cloth_dimx, cloth_dimy]
else:
cloth_dimx, cloth_dimy = config['ClothSize']
self.set_scene(config)
self.action_tool.reset([0., -1., 0.])
pos = pyflex.get_positions().reshape(-1, 4)
pos[:, :3] -= np.mean(pos, axis=0)[:3]
if self.action_mode in ['sawyer', 'franka'
]: # Take care of the table in robot case
pos[:, 1] = 0.57
else:
pos[:, 1] = 0.005
pos[:, 3] = 1
pyflex.set_positions(pos.flatten())
pyflex.set_velocities(np.zeros_like(pos))
pyflex.step()
num_particle = cloth_dimx * cloth_dimy
pickpoint = random.randint(0, num_particle - 1)
curr_pos = pyflex.get_positions()
original_inv_mass = curr_pos[pickpoint * 4 + 3]
curr_pos[
pickpoint * 4 +
3] = 0 # Set the mass of the pickup point to infinity so that it generates enough force to the rest of the cloth
pickpoint_pos = curr_pos[pickpoint * 4:pickpoint * 4 + 3].copy(
) # Pos of the pickup point is fixed to this point
pickpoint_pos[1] += np.random.random(1) * 0.5 + 0.5
pyflex.set_positions(curr_pos)
# Pick up the cloth and wait to stablize
for j in range(0, max_wait_step):
curr_pos = pyflex.get_positions()
curr_vel = pyflex.get_velocities()
curr_pos[pickpoint * 4:pickpoint * 4 + 3] = pickpoint_pos
curr_vel[pickpoint * 3:pickpoint * 3 + 3] = [0, 0, 0]
pyflex.set_positions(curr_pos)
pyflex.set_velocities(curr_vel)
pyflex.step()
if np.alltrue(
np.abs(curr_vel) < stable_vel_threshold) and j > 5:
break
# Drop the cloth and wait to stablize
curr_pos = pyflex.get_positions()
curr_pos[pickpoint * 4 + 3] = original_inv_mass
pyflex.set_positions(curr_pos)
for _ in range(max_wait_step):
pyflex.step()
curr_vel = pyflex.get_velocities()
if np.alltrue(curr_vel < stable_vel_threshold):
break
center_object()
if self.action_mode == 'sphere' or self.action_mode.startswith(
'picker'):
curr_pos = pyflex.get_positions()
self.action_tool.reset(curr_pos[pickpoint * 4:pickpoint * 4 +
3] + [0., 0.2, 0.])
generated_configs.append(deepcopy(config))
generated_states.append(deepcopy(self.get_state()))
self.current_config = config # Needed in _set_to_flatten function
generated_configs[-1]['flatten_area'] = self._set_to_flatten(
) # Record the maximum flatten area
print('config {}: camera params {}, flatten area: {}'.format(
i, config['camera_params'],
generated_configs[-1]['flatten_area']))
return generated_configs, generated_states
def _set_pos(picker_pos, particle_pos):
shape_states = np.array(pyflex.get_shape_states()).reshape(-1, 14)
shape_states[:, 3:6] = shape_states[:, :3]
shape_states[:, :3] = picker_pos
pyflex.set_shape_states(shape_states)
pyflex.set_positions(particle_pos)
def set_scene(self, config, states=None, create_only=False):
'''
Construct the pouring water scence.
'''
# create fluid
super().set_scene(
config
) # do not sample fluid parameters, as it's very likely to generate very strange fluid
# compute glass params
if states is None:
self.set_pouring_glass_params(config["glass"])
self.set_poured_glass_params(config["glass"])
else:
glass_params = states['glass_params']
self.border = glass_params['border']
self.height = glass_params['height']
self.glass_dis_x = glass_params['glass_dis_x']
self.glass_dis_z = glass_params['glass_dis_z']
self.glass_distance = glass_params['glass_distance']
self.poured_border = glass_params['poured_border']
self.poured_height = glass_params['poured_height']
self.poured_glass_dis_x = glass_params['poured_glass_dis_x']
self.poured_glass_dis_z = glass_params['poured_glass_dis_z']
self.glass_params = glass_params
# create pouring glass & poured glass
self.create_glass(self.glass_dis_x, self.glass_dis_z, self.height,
self.border)
self.create_glass(self.poured_glass_dis_x, self.poured_glass_dis_z,
self.poured_height, self.poured_border)
# move pouring glass to be at ground
self.glass_states = self.init_glass_state(self.x_center, 0,
self.glass_dis_x,
self.glass_dis_z,
self.height, self.border)
# move poured glass to be at ground
self.poured_glass_states = self.init_glass_state(
self.x_center + self.glass_distance, 0, self.poured_glass_dis_x,
self.poured_glass_dis_z, self.poured_height, self.poured_border)
self.set_shape_states(self.glass_states, self.poured_glass_states)
# record glass floor center x, y, and rotation
self.glass_x = self.x_center
if self.action_mode == 'rotation_bottom':
self.glass_y = 0
elif self.action_mode == 'rotation_top':
self.glass_y = 0.5 * self.border + self.height
self.glass_rotation = 0
# only create the glass and water, without setting their states
# this is only used in the pourwater amount env.
if create_only:
return
# no cached init states passed in
if states is None:
fluid_pos = np.ones((self.particle_num, self.dim_position))
# move water all inside the glass
fluid_radius = self.fluid_params['radius'] * self.fluid_params[
'rest_dis_coef']
fluid_dis = np.array(
[1.0 * fluid_radius, fluid_radius * 0.5, 1.0 * fluid_radius])
lower_x = self.glass_params['glass_x_center'] - self.glass_params[
'glass_dis_x'] / 2. + self.glass_params['border']
lower_z = -self.glass_params[
'glass_dis_z'] / 2 + self.glass_params['border']
lower_y = self.glass_params['border']
if self.action_mode in ['sawyer', 'franka']:
lower_y += 0.56 # NOTE: robotics table
lower = np.array([lower_x, lower_y, lower_z])
cnt = 0
rx = int(self.fluid_params['dim_x'] * 1)
ry = int(self.fluid_params['dim_y'] * 1)
rz = int(self.fluid_params['dim_z'] / 1)
for x in range(rx):
for y in range(ry):
for z in range(rz):
fluid_pos[cnt][:3] = lower + np.array(
[x, y, z]) * fluid_dis # + np.random.rand() * 0.01
cnt += 1
pyflex.set_positions(fluid_pos)
print("stablize water!")
for _ in range(100):
pyflex.step()
state_dic = self.get_state()
water_state = state_dic['particle_pos'].reshape(
(-1, self.dim_position))
in_glass = self.in_glass(water_state, self.glass_states,
self.border, self.height)
not_in_glass = 1 - in_glass
not_total_num = np.sum(not_in_glass)
while not_total_num > 0:
max_height_now = np.max(water_state[:, 1])
fluid_dis = np.array(
[1.0 * fluid_radius, fluid_radius * 1, 1.0 * fluid_radius])
lower_x = self.glass_params[
'glass_x_center'] - self.glass_params['glass_dis_x'] / 4
lower_z = -self.glass_params['glass_dis_z'] / 4
lower_y = max_height_now
lower = np.array([lower_x, lower_y, lower_z])
cnt = 0
dim_x = config['fluid']['dim_x']
dim_z = config['fluid']['dim_z']
for w_idx in range(len(water_state)):
if not in_glass[w_idx]:
water_state[w_idx][:3] = lower + fluid_dis * np.array([
cnt % dim_x, cnt // (dim_x * dim_z),
(cnt // dim_x) % dim_z
])
cnt += 1
pyflex.set_positions(water_state)
for _ in range(40):
pyflex.step()
state_dic = self.get_state()
water_state = state_dic['particle_pos'].reshape(
(-1, self.dim_position))
in_glass = self.in_glass(water_state, self.glass_states,
self.border, self.height)
not_in_glass = 1 - in_glass
not_total_num = np.sum(not_in_glass)
for _ in range(30):
pyflex.step()
else: # set to passed-in cached init states
self.set_state(states)
def set_scene(self, config, states=None):
'''
Construct the passing water scence.
'''
# create fluid
super().set_scene(
config
) # do not sample fluid parameters, as it's very likely to generate very strange fluid
# compute glass params
if states is None:
self.set_glass_params(config["glass"])
else:
glass_params = states['glass_params']
self.border = glass_params['border']
self.height = glass_params['height']
self.glass_dis_x = glass_params['glass_dis_x']
self.glass_dis_z = glass_params['glass_dis_z']
self.glass_params = glass_params
# create glass
self.create_glass(self.glass_dis_x, self.glass_dis_z, self.height,
self.border)
# move glass to be at ground or on the table
self.glass_states = self.init_glass_state(self.x_center, 0,
self.glass_dis_x,
self.glass_dis_z,
self.height, self.border)
pyflex.set_shape_states(self.glass_states)
# record glass floor center x
self.glass_x = self.x_center
# no cached init states passed in
if states is None:
fluid_pos = np.ones((self.particle_num, self.dim_position))
# move water all inside the glass
fluid_radius = self.fluid_params['radius'] * self.fluid_params[
'rest_dis_coef']
fluid_dis = np.array(
[1.0 * fluid_radius, fluid_radius * 0.5, 1.0 * fluid_radius])
lower_x = self.glass_params['glass_x_center'] - self.glass_params[
'glass_dis_x'] / 2. + self.glass_params['border']
lower_z = -self.glass_params[
'glass_dis_z'] / 2 + self.glass_params['border']
lower_y = self.glass_params['border']
lower = np.array([lower_x, lower_y, lower_z])
cnt = 0
rx = int(self.fluid_params['dim_x'] * 1)
ry = int(self.fluid_params['dim_y'] * 1)
rz = int(self.fluid_params['dim_z'] / 1)
for x in range(rx):
for y in range(ry):
for z in range(rz):
fluid_pos[cnt][:3] = lower + np.array([x, y, z
]) * fluid_dis
cnt += 1
pyflex.set_positions(fluid_pos)
print("stablize water!")
for _ in range(100):
pyflex.step()
state_dic = self.get_state()
water_state = state_dic['particle_pos'].reshape(
(-1, self.dim_position))
in_glass = self.in_glass(water_state,
self.glass_states,
self.border,
self.height,
return_sum=False)
not_in_glass = 1 - in_glass
not_total_num = np.sum(not_in_glass)
while not_total_num > 0:
max_height_now = np.max(water_state[:, 1])
fluid_dis = np.array(
[1.0 * fluid_radius, fluid_radius * 1, 1.0 * fluid_radius])
lower_x = self.glass_params[
'glass_x_center'] - self.glass_params['glass_dis_x'] / 4
lower_z = -self.glass_params['glass_dis_z'] / 4
lower_y = max_height_now
lower = np.array([lower_x, lower_y, lower_z])
cnt = 0
dim_x = config['fluid']['dim_x']
dim_z = config['fluid']['dim_z']
for w_idx in range(len(water_state)):
if not in_glass[w_idx]:
water_state[w_idx][:3] = lower + fluid_dis * np.array([
cnt % dim_x, cnt // (dim_x * dim_z),
(cnt // dim_x) % dim_z
])
cnt += 1
pyflex.set_positions(water_state)
for _ in range(100):
pyflex.step()
state_dic = self.get_state()
water_state = state_dic['particle_pos'].reshape(
(-1, self.dim_position))
in_glass = self.in_glass(water_state,
self.glass_states,
self.border,
self.height,
return_sum=False)
not_in_glass = 1 - in_glass
not_total_num = np.sum(not_in_glass)
for _ in range(30):
pyflex.step()
else: # set to passed-in cached init states
self.set_state(states)
