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)