-
Notifications
You must be signed in to change notification settings - Fork 3
/
run_simulation_diff_nogui.py
68 lines (57 loc) · 2.15 KB
/
run_simulation_diff_nogui.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
import os
import argparse
import taichi as ti
import numpy as np
from config_builder import SimConfig
from particle_system_diff import ParticleSystem
from DFSPH_diff import DFSPHSolver
ti.init(arch=ti.gpu, device_memory_fraction=0.5, debug=True)
def build_solver(ps: ParticleSystem):
solver_type = ps.cfg.get_cfg("simulationMethod")
# if solver_type == 0:
# return WCSPHSolver(ps)
# elif solver_type == 4:
if solver_type == 4:
return DFSPHSolver(ps)
else:
raise NotImplementedError(f"Solver type {solver_type} has not been implemented.")
if __name__ == "__main__":
parser = argparse.ArgumentParser(description='SPH Taichi')
parser.add_argument('--scene_file',
default='',
help='scene file')
args = parser.parse_args()
scene_path = args.scene_file
config = SimConfig(scene_file_path=scene_path)
scene_name = scene_path.split("/")[-1].split(".")[0]
substeps = config.get_cfg("numberOfStepsPerRenderUpdate")
output_frames = config.get_cfg("exportFrame")
output_interval = int(0.016 / config.get_cfg("timeStepSize"))
output_ply = config.get_cfg("exportPly")
output_obj = config.get_cfg("exportObj")
series_prefix = "{}_output/particle_object_{}.ply".format(scene_name, "{}")
if output_frames:
os.makedirs(f"{scene_name}_output_img", exist_ok=True)
if output_ply:
os.makedirs(f"{scene_name}_output", exist_ok=True)
ps = ParticleSystem(config, GGUI=True)
solver = build_solver(ps)
solver.initialize()
cnt_frame = 0
losses = []
while True:
# TODO:
# the code cannot work this way
# some replaced function grads should be modified
with ti.ad.Tape(loss=ps.loss, validation=True):
solver.initialize_from_restart()
while not solver.end(cnt_frame):
print(cnt_frame)
solver.step(cnt_frame)
cnt_frame += 1
print("finish")
solver.compute_loss(ps.steps - 1)
current_loss = ps.loss[None]
losses.append(current_loss)
print("loss: ", current_loss)
solver.update()