Skip to content

Commit 31a7c1f

Browse files
chahak13chahak13
committed
Add control flow for USF, USL, MUSL
1 parent 6b54359 commit 31a7c1f

3 files changed

Lines changed: 83 additions & 78 deletions

File tree

‎diffmpm/mesh.py‎

Lines changed: 35 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -426,7 +426,6 @@ def _transfer_node_force_vel_par(self, dt):
426426
)
427427
)
428428

429-
# @jit
430429
def _update_par_pos_node_mom(self, dt):
431430
"""
432431
Update particle position based on nodal momentum.
@@ -546,7 +545,6 @@ def step(pid, args):
546545
0, len(self.particles), step, args
547546
)
548547

549-
# @jit
550548
def _update_node_fint_par_mass(self):
551549
r"""
552550
Update the nodal internal force based on particle mass.
@@ -596,7 +594,6 @@ def step(pid, args):
596594
0, len(self.particles), step, args
597595
)
598596

599-
# @jit
600597
def _update_node_fext_par_fext(self):
601598
r"""
602599
Update the nodal external force based on particle f_ext.
@@ -636,15 +633,10 @@ def solve(self, nsteps=100, mpm_scheme="USF", **kwargs):
636633
"velocity": [],
637634
}
638635
for _ in tqdm(range(nsteps)):
639-
# print("="*30)
640636
self._update_particle_natural_coords()
641-
# print(f"P.xi: {self.particles.xi}")
642637
self._update_particle_element_ids()
643-
# print(f"P.eid: {self.particles.element_ids}")
644638
self._update_node_momentum_par_vel()
645-
# print(f"P.v: {self.particles.velocity}")
646639
self._update_node_mass_par_mass()
647-
# print(f"P.m: {self.particles.mass}")
648640
self._update_nodes_bc_mom_vel()
649641
if mpm_scheme == "USF":
650642
self._update_nodes_mom_vel()
@@ -674,7 +666,6 @@ def solve(self, nsteps=100, mpm_scheme="USF", **kwargs):
674666
result = {k: jnp.asarray(v) for k, v in result.items()}
675667
return result
676668

677-
# @partial(jit, static_argnames=["nsteps", "mpm_scheme"])
678669
def solve_jit(self, nsteps=100, mpm_scheme=0, **kwargs):
679670
"""
680671
Solve the mesh using explicit scheme (for now).
@@ -691,6 +682,8 @@ def solve_jit(self, nsteps=100, mpm_scheme=0, **kwargs):
691682
"position": jnp.zeros((nsteps, self.particles.nparticles)),
692683
"velocity": jnp.zeros((nsteps, self.particles.nparticles)),
693684
}
685+
mpm_scheme_dict = {"USF": 0, "USL": 1, "MUSL": 2}
686+
mpm_scheme = mpm_scheme_dict[mpm_scheme]
694687

695688
@loop_tqdm(nsteps)
696689
def step(i, data):
@@ -703,19 +696,20 @@ def step(i, data):
703696

704697
# debug.breakpoint()
705698
# if mpm_scheme == 0:
706-
# def f(data):
707-
# s, dt = data
708-
# s._update_nodes_mom_vel()
709-
# s._update_particle_strain(dt)
710-
# s._update_par_vol_density()
711-
# s._update_particle_stress()
712-
# return None
713-
714-
# lax.cond(mpm_scheme == 0, f, lambda data: None, (self, dt))
715-
self._update_nodes_mom_vel()
716-
self._update_particle_strain(dt)
717-
self._update_par_vol_density()
718-
self._update_particle_stress()
699+
def f(data):
700+
# debug.breakpoint()
701+
s, dt = data
702+
s._update_nodes_mom_vel()
703+
s._update_particle_strain(dt)
704+
s._update_par_vol_density()
705+
s._update_particle_stress()
706+
return s
707+
708+
# TODO: Confirm that it is fine doing this since
709+
# it sets a value to self.
710+
self = lax.cond(
711+
mpm_scheme == 0, f, lambda data: data[0], (self, dt)
712+
)
719713

720714
self._update_node_fint_par_mass()
721715
self._update_node_fext_par_fext()
@@ -724,22 +718,27 @@ def step(i, data):
724718
self._transfer_node_force_vel_par(dt)
725719
self._update_par_pos_node_mom(dt)
726720

727-
# def f(s):
728-
# s._update_node_momentum_par_vel()
729-
# s._update_nodes_bc_mom_vel()
730-
# return None
721+
def f(s):
722+
s._update_node_momentum_par_vel()
723+
s._update_nodes_bc_mom_vel()
724+
return s
731725

732-
# lax.cond(mpm_scheme == 2, f, lambda s: None, self)
726+
self = lax.cond(mpm_scheme == 2, f, lambda s: s, self)
733727

734-
# def f(data):
735-
# s, dt = data
736-
# s._update_nodes_mom_vel()
737-
# s._update_particle_strain(dt)
738-
# s._update_par_vol_density()
739-
# s._update_particle_stress()
740-
# return None
728+
def f(data):
729+
s, dt = data
730+
s._update_nodes_mom_vel()
731+
s._update_particle_strain(dt)
732+
s._update_par_vol_density()
733+
s._update_particle_stress()
734+
return s
741735

742-
# lax.cond(mpm_scheme in {1, 2}, f, lambda data: None, (self, dt))
736+
self = lax.cond(
737+
mpm_scheme == 1, f, lambda data: data[0], (self, dt)
738+
)
739+
self = lax.cond(
740+
mpm_scheme == 2, f, lambda data: data[0], (self, dt)
741+
)
743742
self.nodes.reset_values()
744743
result["position"] = (
745744
result["position"].at[i, :].set(self.particles.x)
@@ -750,7 +749,7 @@ def step(i, data):
750749
return (self, mpm_scheme, dt, result)
751750

752751
_, _, _, result = lax.fori_loop(
753-
0, nsteps, step, (self, 0, kwargs["dt"], result)
752+
0, nsteps, step, (self, mpm_scheme, kwargs["dt"], result)
754753
)
755754
result = {k: jnp.asarray(v) for k, v in result.items()}
756755
return result

‎jax_port.org‎

Lines changed: 35 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -824,7 +824,6 @@ class Mesh1D:
824824
)
825825
)
826826

827-
# @jit
828827
def _update_par_pos_node_mom(self, dt):
829828
"""
830829
Update particle position based on nodal momentum.
@@ -944,7 +943,6 @@ class Mesh1D:
944943
0, len(self.particles), step, args
945944
)
946945

947-
# @jit
948946
def _update_node_fint_par_mass(self):
949947
r"""
950948
Update the nodal internal force based on particle mass.
@@ -994,7 +992,6 @@ class Mesh1D:
994992
0, len(self.particles), step, args
995993
)
996994

997-
# @jit
998995
def _update_node_fext_par_fext(self):
999996
r"""
1000997
Update the nodal external force based on particle f_ext.
@@ -1034,15 +1031,10 @@ class Mesh1D:
10341031
"velocity": [],
10351032
}
10361033
for _ in tqdm(range(nsteps)):
1037-
# print("="*30)
10381034
self._update_particle_natural_coords()
1039-
# print(f"P.xi: {self.particles.xi}")
10401035
self._update_particle_element_ids()
1041-
# print(f"P.eid: {self.particles.element_ids}")
10421036
self._update_node_momentum_par_vel()
1043-
# print(f"P.v: {self.particles.velocity}")
10441037
self._update_node_mass_par_mass()
1045-
# print(f"P.m: {self.particles.mass}")
10461038
self._update_nodes_bc_mom_vel()
10471039
if mpm_scheme == "USF":
10481040
self._update_nodes_mom_vel()
@@ -1072,7 +1064,6 @@ class Mesh1D:
10721064
result = {k: jnp.asarray(v) for k, v in result.items()}
10731065
return result
10741066

1075-
# @partial(jit, static_argnames=["nsteps", "mpm_scheme"])
10761067
def solve_jit(self, nsteps=100, mpm_scheme=0, **kwargs):
10771068
"""
10781069
Solve the mesh using explicit scheme (for now).
@@ -1089,6 +1080,8 @@ class Mesh1D:
10891080
"position": jnp.zeros((nsteps, self.particles.nparticles)),
10901081
"velocity": jnp.zeros((nsteps, self.particles.nparticles)),
10911082
}
1083+
mpm_scheme_dict = {"USF": 0, "USL": 1, "MUSL": 2}
1084+
mpm_scheme = mpm_scheme_dict[mpm_scheme]
10921085

10931086
@loop_tqdm(nsteps)
10941087
def step(i, data):
@@ -1101,19 +1094,20 @@ class Mesh1D:
11011094

11021095
# debug.breakpoint()
11031096
# if mpm_scheme == 0:
1104-
# def f(data):
1105-
# s, dt = data
1106-
# s._update_nodes_mom_vel()
1107-
# s._update_particle_strain(dt)
1108-
# s._update_par_vol_density()
1109-
# s._update_particle_stress()
1110-
# return None
1111-
1112-
# lax.cond(mpm_scheme == 0, f, lambda data: None, (self, dt))
1113-
self._update_nodes_mom_vel()
1114-
self._update_particle_strain(dt)
1115-
self._update_par_vol_density()
1116-
self._update_particle_stress()
1097+
def f(data):
1098+
# debug.breakpoint()
1099+
s, dt = data
1100+
s._update_nodes_mom_vel()
1101+
s._update_particle_strain(dt)
1102+
s._update_par_vol_density()
1103+
s._update_particle_stress()
1104+
return s
1105+
1106+
# TODO: Confirm that it is fine doing this since
1107+
# it sets a value to self.
1108+
self = lax.cond(
1109+
mpm_scheme == 0, f, lambda data: data[0], (self, dt)
1110+
)
11171111

11181112
self._update_node_fint_par_mass()
11191113
self._update_node_fext_par_fext()
@@ -1122,22 +1116,27 @@ class Mesh1D:
11221116
self._transfer_node_force_vel_par(dt)
11231117
self._update_par_pos_node_mom(dt)
11241118

1125-
# def f(s):
1126-
# s._update_node_momentum_par_vel()
1127-
# s._update_nodes_bc_mom_vel()
1128-
# return None
1119+
def f(s):
1120+
s._update_node_momentum_par_vel()
1121+
s._update_nodes_bc_mom_vel()
1122+
return s
11291123

1130-
# lax.cond(mpm_scheme == 2, f, lambda s: None, self)
1124+
self = lax.cond(mpm_scheme == 2, f, lambda s: s, self)
11311125

1132-
# def f(data):
1133-
# s, dt = data
1134-
# s._update_nodes_mom_vel()
1135-
# s._update_particle_strain(dt)
1136-
# s._update_par_vol_density()
1137-
# s._update_particle_stress()
1138-
# return None
1126+
def f(data):
1127+
s, dt = data
1128+
s._update_nodes_mom_vel()
1129+
s._update_particle_strain(dt)
1130+
s._update_par_vol_density()
1131+
s._update_particle_stress()
1132+
return s
11391133

1140-
# lax.cond(mpm_scheme in {1, 2}, f, lambda data: None, (self, dt))
1134+
self = lax.cond(
1135+
mpm_scheme == 1, f, lambda data: data[0], (self, dt)
1136+
)
1137+
self = lax.cond(
1138+
mpm_scheme == 2, f, lambda data: data[0], (self, dt)
1139+
)
11411140
self.nodes.reset_values()
11421141
result["position"] = (
11431142
result["position"].at[i, :].set(self.particles.x)
@@ -1148,7 +1147,7 @@ class Mesh1D:
11481147
return (self, mpm_scheme, dt, result)
11491148

11501149
_, _, _, result = lax.fori_loop(
1151-
0, nsteps, step, (self, 0, kwargs["dt"], result)
1150+
0, nsteps, step, (self, mpm_scheme, kwargs["dt"], result)
11521151
)
11531152
result = {k: jnp.asarray(v) for k, v in result.items()}
11541153
return result

‎test/test_mesh.py‎

Lines changed: 13 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -258,10 +258,10 @@ def test_analytical_solve():
258258
mesh.set_particle_velocity(velocity)
259259
mesh2.set_particle_velocity(velocity)
260260
dt = 0.01
261-
duration = 25
261+
duration = 1
262262
nsteps = int(duration / dt)
263-
r = mesh2.solve(dt=dt, nsteps=nsteps, mpm_scheme="USF")
264-
result = mesh.solve_jit(dt=dt, nsteps=nsteps, mpm_scheme="USF")
263+
r = mesh2.solve(dt=dt, nsteps=nsteps, mpm_scheme="USL")
264+
result = mesh.solve_jit(dt=dt, nsteps=nsteps, mpm_scheme="USL")
265265

266266
def analytical_vibration(E, rho, v0, x_loc, duration, dt, L):
267267
t = jnp.arange(0, duration, dt)
@@ -276,12 +276,19 @@ def analytical_vibration(E, rho, v0, x_loc, duration, dt, L):
276276
ax.plot(
277277
t,
278278
result["velocity"],
279-
"k+-",
279+
"k+--",
280280
linewidth=1,
281-
markersize=4,
281+
markersize=6,
282282
label="diffmpm-jit",
283283
)
284-
ax.plot(t, r["velocity"], "go-", linewidth=1, markersize=4, label="diffmpm")
284+
ax.plot(
285+
t,
286+
r["velocity"],
287+
"gx--",
288+
linewidth=1,
289+
markersize=4,
290+
label="diffmpm",
291+
)
285292
ax.grid()
286293
ax.legend()
287294
ax.set_xlabel("Time (s)")

0 commit comments

Comments
 (0)