diff --git a/python/taichi/examples/simulation/implicit_mass_spring.py b/python/taichi/examples/simulation/implicit_mass_spring.py index d2ba9a335731d..adf47aa574c2d 100644 --- a/python/taichi/examples/simulation/implicit_mass_spring.py +++ b/python/taichi/examples/simulation/implicit_mass_spring.py @@ -237,20 +237,32 @@ def displayGGUI(self, canvas, radius=0.01, color=(1.0, 1.0, 1.0)): def main(): - ti.init(arch=ti.cpu) - h = 0.01 - cloth = Cloth(N=5) - - pause = False parser = argparse.ArgumentParser() parser.add_argument('-g', '--use-ggui', action='store_true', help='Display with GGUI') + parser.add_argument('-a', + '--arch', + required=False, + default="cpu", + dest='arch', + type=str, + help='The arch (backend) to run this example on') args, unknowns = parser.parse_known_args() - use_ggui = False - use_ggui = args.use_ggui + arch = args.arch + if arch in ["x64", "cpu", "arm64"]: + ti.init(arch=ti.cpu) + elif arch in ["cuda", "gpu"]: + ti.init(arch=ti.cuda) + else: + raise ValueError('Only CPU and CUDA backends are supported for now.') + h = 0.01 + pause = False + cloth = Cloth(N=5) + + use_ggui = args.use_ggui if not use_ggui: gui = ti.GUI('Implicit Mass Spring System', res=(500, 500)) while gui.running: diff --git a/python/taichi/examples/simulation/stable_fluid.py b/python/taichi/examples/simulation/stable_fluid.py index ea20cb177209c..ddcdcdc2e87ba 100644 --- a/python/taichi/examples/simulation/stable_fluid.py +++ b/python/taichi/examples/simulation/stable_fluid.py @@ -18,9 +18,16 @@ '--use-sp-mat', action='store_true', help='Solve Poisson\'s equation by using a sparse matrix') +parser.add_argument('-a', + '--arch', + required=False, + default="cpu", + dest='arch', + type=str, + help='The arch (backend) to run this example on') args, unknowns = parser.parse_known_args() -res = 512 +res = 256 dt = 0.03 p_jacobi_iters = 500 # 40 for a quicker but less accurate result f_strength = 10000.0 @@ -32,12 +39,17 @@ debug = False use_sparse_matrix = args.use_sp_mat +arch = args.arch +if arch in ["x64", "cpu", "arm64"]: + ti.init(arch=ti.cpu) +elif arch in ["cuda", "gpu"]: + ti.init(arch=ti.cuda) +else: + raise ValueError('Only CPU and CUDA backends are supported for now.') if use_sparse_matrix: - ti.init(arch=ti.x64) print('Using sparse matrix') else: - ti.init(arch=ti.gpu) print('Using jacobi iteration') _velocities = ti.Vector.field(2, float, shape=(res, res)) @@ -86,7 +98,7 @@ def fill_laplacian_matrix(A: ti.types.sparse_matrix_builder()): N = res * res K = ti.linalg.SparseMatrixBuilder(N, N, max_num_triplets=N * 6) - F_b = ti.field(ti.f32, shape=N) + F_b = ti.ndarray(ti.f32, shape=N) fill_laplacian_matrix(K) L = K.build() @@ -236,7 +248,7 @@ def enhance_vorticity(vf: ti.template(), cf: ti.template()): @ti.kernel -def copy_divergence(div_in: ti.template(), div_out: ti.template()): +def copy_divergence(div_in: ti.template(), div_out: ti.types.ndarray()): for I in ti.grouped(div_in): div_out[I[0] * res + I[1]] = -div_in[I]