Example #1
0
def main():
    parser = parse_mpm()
    args = parser.parse_args()
    simulation_trajectories = []

    if args.material == 'water':
        material = MPMSolver.material_water
    elif args.material == 'elastic':
        material = MPMSolver.material_elastic
    elif args.material == 'snow':
        material = MPMSolver.material_snow
    elif args.material == 'sand':
        material = MPMSolver.material_sand

    for sim_i in range(args.num_sims):
        ti.init(arch=ti.cuda)  # Try to run on GPU
        gui = ti.GUI("Taichi MLS-MPM-99", res=512, background_color=0x112F41)

        mpm = MPMSolver(res=(128, 128))
        mpm.E = args.young
        # init condition
        n_samples = args.ncubes * 10

        lower_corner_candidates = np.random.uniform(0.2,
                                                    0.8,
                                                    size=(n_samples, 2))
        cube_size_candidates = np.random.uniform(0.06,
                                                 0.15,
                                                 size=(n_samples, 2))

        lower_corners = []
        cube_sizes = []
        num_cubes = args.ncubes
        for s in range(1, n_samples):
            if not cube_overlap(
                    lower_corner_candidates[:s], cube_size_candidates[:s],
                    lower_corner_candidates[s], cube_size_candidates[s]):
                lower_corners.append(lower_corner_candidates[s])
                cube_sizes.append(cube_size_candidates[s])
            if len(lower_corners) == num_cubes:
                break

        for i in range(len(lower_corners)):
            mpm.add_cube(lower_corner=lower_corners[i],
                         cube_size=cube_sizes[i],
                         material=material)

        simulation_trajectories.append(simulate(mpm, gui, args))
        #rollout(simulation_trajectories[0], gui, args.gui_particle_radius)
        #cluster_scene(simulation_trajectories[-1][0])

        if (sim_i + 1) % 1 == 0:
            print('Simulated {} trajectories.'.format(sim_i + 1))

        data_path = os.path.join('../data/', args.dataset_name)
        if not os.path.exists(data_path):
            os.makedirs(data_path)

        if (sim_i + 1) % 100 == 0 or sim_i + 1 == args.num_sims:
            print(len(simulation_trajectories))
            with open(
                    '../data/{}/{}_{}.pkl'.format(args.dataset_name,
                                                  args.material, sim_i // 100),
                    'wb') as f:
                pkl.dump(simulation_trajectories, f)
            simulation_trajectories = []