Skip to content

Commit

Permalink
clean up swarm code
Browse files Browse the repository at this point in the history
  • Loading branch information
@lmoresi
lmoresi committed Apr 4, 2024
1 parent d67f6fd commit 0c7107d
Showing 1 changed file with 0 additions and 178 deletions.
178 changes: 0 additions & 178 deletions src/underworld3/swarm.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1932,181 +1932,3 @@ def __init__(
self._X0 = nX0

return

# @timing.routine_timer_decorator
# def estimate_dt(self, V_fn):
# """
# Calculates an appropriate advective timestep for the given
# mesh and velocity configuration.
# """
# # we'll want to do this on an element by element basis
# # for more general mesh

# # first let's extract a max global velocity magnitude
# import math

# with self.access():
# vel = uw.function.evalf(V_fn, self.particle_coordinates.data)
# magvel_squared = vel[:, 0] ** 2 + vel[:, 1] ** 2
# if self.mesh.dim == 3:
# magvel_squared += vel[:, 2] ** 2

# max_magvel = math.sqrt(magvel_squared.max())

# from mpi4py import MPI

# comm = MPI.COMM_WORLD
# max_magvel_glob = comm.allreduce(max_magvel, op=MPI.MAX)

# min_dx = self.mesh.get_min_radius()

# # The assumption should be that we cross one or two elements (4 radii), not more,
# # in a single step (order 2, means one element per half-step or something
# # that we can broadly interpret that way)

# if max_magvel_glob != 0.0:
# return 4.0 * min_dx / max_magvel_glob
# else:
# return None

@timing.routine_timer_decorator
def advection(
self,
V_fn,
delta_t,
order=2,
corrector=False,
restore_points_to_domain_func=None,
evalf=False,
):
# X0 holds the particle location at the start of advection
# This is needed because the particles may be migrated off-proc
# during timestepping.

dt_limit = self.estimate_dt(V_fn)

if dt_limit is not None:
substeps = int(max(1, abs(delta_t) // dt_limit))
else:
substeps = 1

X0 = self._X0

# Use current velocity to estimate where the particles would have
# landed in an implicit step.

# ? how does this interact with the particle restoration function ?

V_fn_matrix = self.mesh.vector.to_matrix(V_fn)

with self.access(X0):
X0.data[...] = self.data[...]
self._X0_uninitialised = False

# Wrap this whole thing in sub-stepping loop

for step in range(0, substeps):

# Mid point algorithm (2nd order)
if order == 2:
with self.access(self.particle_coordinates):
v_at_Vpts = np.zeros_like(self.data)

if evalf:
for d in range(self.dim):
v_at_Vpts[:, d] = uw.function.evalf(
V_fn_matrix[d], self.data
).reshape(-1)
else:
for d in range(self.dim):
v_at_Vpts[:, d] = uw.function.evaluate(
V_fn_matrix[d], self.data
).reshape(-1)

bc_mask_array = np.rint(
uw.function.evalf(_bc_mask_fn, self.data)
).reshape(-1, 1)

mid_pt_coords = (
self.data[...].copy() + 0.5 * delta_t * v_at_Vpts / substeps
)

# validate_coords to ensure they live within the domain (or there will be trouble)

if restore_points_to_domain_func is not None:
mid_pt_coords = restore_points_to_domain_func(mid_pt_coords)

self.data[...] = mid_pt_coords[...]

del mid_pt_coords

## Let the swarm be updated, and then move the rest of the way

with self.access(self.particle_coordinates):
v_at_Vpts = np.zeros_like(self.data)

if evalf:
for d in range(self.dim):
v_at_Vpts[:, d] = uw.function.evalf(
V_fn_matrix[d], self.data
).reshape(-1)
else:
for d in range(self.dim):
v_at_Vpts[:, d] = uw.function.evaluate(
V_fn_matrix[d], self.data
).reshape(-1)

# if (uw.mpi.rank == 0):
# print("Re-launch from X0", flush=True)

bc_mask_array = np.rint(
uw.function.evalf(_bc_mask_fn, self.data)
).reshape(-1, 1)
new_coords = X0.data[...].copy() + delta_t * v_at_Vpts / substeps

# validate_coords to ensure they live within the domain (or there will be trouble)
if restore_points_to_domain_func is not None:
new_coords = restore_points_to_domain_func(new_coords)

self.data[...] = new_coords[...]

del new_coords
del v_at_Vpts

# Previous position algorithm (cf above) - we use the previous step as the
# launch point using the current velocity field. This gives a correction to the previous
# landing point.

# assumes X0 is stored from the previous step ... midpoint is needed in the first step

# forward Euler (1st order)
else:
with self.access(self.particle_coordinates):
v_at_Vpts = np.zeros_like(self.data)

if evalf:
for d in range(self.dim):
v_at_Vpts[:, d] = uw.function.evalf(
V_fn_matrix[d], self.data
).reshape(-1)
else:
for d in range(self.dim):
v_at_Vpts[:, d] = uw.function.evaluate(
V_fn_matrix[d], self.data
).reshape(-1)

bc_mask_array = np.rint(
uw.function.evalf(_bc_mask_fn, self.data)
).reshape(-1, 1)
new_coords = self.data + delta_t * v_at_Vpts / substeps

# validate_coords to ensure they live within the domain (or there will be trouble)

if restore_points_to_domain_func is not None:
new_coords = restore_points_to_domain_func(new_coords)

self.data[...] = new_coords[...].copy()

# print(f"Substep - {step}", flush=True)

# End substepping loop

0 comments on commit 0c7107d

Please sign in to comment.