import numpy as np
from scipy.spatial import cKDTree
class DPMSolverSplit2D:
def __init__(
self,
nx=40,
ny=40,
lx=40.0,
ly=40.0,
dt=0.05,
search_radius_factor=2.5,
cos_theta_threshold=-0.1,
velocity_type="quasilinear_transport",
omega_0=0.04,
):
"""Initialize solver parameters and particle grid."""
self.nx, self.ny, self.lx, self.ly = nx, ny, lx, ly
self.dt = dt
self.search_radius_factor = search_radius_factor
self.cos_theta_threshold = cos_theta_threshold
self.velocity_type = velocity_type
self.omega_0 = omega_0
valid_types = [
"quasilinear_transport",
"linear_transport",
"linear_rotation",
"quasilinear_rotation",
"rotating_density_blob",
"rankine_vortex",
"quasilinear_rotation_pressure",
]
if velocity_type not in valid_types:
raise ValueError(f"velocity_type must be one of {valid_types}")
self.x0, self.y0 = lx / 2.0, ly / 2.0
self.R = min(lx, ly) / 5.0
self._initialize_particles()
self._compute_initial_masses()
self.initial_total_mass = self.compute_total_mass()
print(f"Initialized DPM Solver ({velocity_type}) with {nx*ny} particles")
def _initialize_particles(self):
"""Uniform grid initialization with paraboloid height."""
xs = np.linspace(0, self.lx, self.nx)
ys = np.linspace(0, self.ly, self.ny)
X, Y = np.meshgrid(xs, ys, indexing="ij")
self.X, self.Y = X.ravel(), Y.ravel()
self.N = len(self.X)
r2 = (self.X - self.x0) ** 2 + (self.Y - self.y0) ** 2
self.H = np.where(r2 <= self.R**2, 1.0 - r2 / self.R**2, 0.0)
self.H_initial = self.H.copy()
self.dx = self.lx / (self.nx - 1) if self.nx > 1 else self.lx
self.dy = self.ly / (self.ny - 1) if self.ny > 1 else self.ly
if hasattr(self, "velocity_type") and self.velocity_type == "rankine_vortex":
r = np.sqrt(r2)
r_core = self.R / 2.0
denesity_factor = np.where(
r < r_core,
np.exp(-0.5 * (self.omega_0 * r) ** 2),
np.exp(-0.5 * (self.omega_0 * r_core**2 / r) ** 2),
)
self.H = self.H * (0.3 + 0.7 * denesity_factor)
self.H_initial = self.H.copy()
def velocity_field(self, X, Y, H):
"""Return velocity components (U,V) for given field type."""
if self.velocity_type == "quasilinear_transport":
U = 0.5 * H
V = 0.5 * H
elif self.velocity_type == "linear_transport":
U = 0.5 * np.ones_like(H)
V = 0.5 * np.ones_like(H)
elif self.velocity_type == "linear_rotation":
dx, dy = X - self.x0, Y - self.y0
U, V = -dy * self.omega_0, dx * self.omega_0
elif self.velocity_type == "quasilinear_rotation":
dx, dy = X - self.x0, Y - self.y0
omega = self.omega_0 * H
U, V = -dy * omega, dx * omega
elif self.velocity_type == "rotating_density_blob":
dx, dy = X - self.x0, Y - self.y0
r = np.hypot(dx, dy) + 1e-10
omega = self.omega_0 * (1.0 + H)
U_rot = -dy * omega
V_rot = dx * omega
dH_dr = -2.0 * r / self.R**2 * H
strength = 0.05
U = U_rot + strength * dH_dr * (dx / r)
V = V_rot + strength * dH_dr * (dy / r)
elif self.velocity_type == "rankine_vortex":
dx, dy = X - self.x0, Y - self.y0
r = np.hypot(dx, dy) + 1e-10
r_core = self.R / 2.0
v_theta = np.where(
r < r_core, self.omega_0 * r, self.omega_0 * r_core**2 / r
)
U = -dy / r * v_theta
V = dx / r * v_theta
elif self.velocity_type == "quasilinear_rotation_pressure":
dx, dy = X - self.x0, Y - self.y0
r = np.hypot(dx, dy) + 1e-12
omega = self.omega_0 * H
U_rot = -dy * omega
V_rot = dx * omega
alpha = 0.05 # давление
# радиальный градиент H (аналитически, как у тебя в blob)
dH_dr = -2.0 * r / self.R**2 * H
U = U_rot - alpha * dH_dr * (dx / r)
V = V_rot - alpha * dH_dr * (dy / r)
return U, V
def _compute_initial_masses(self):
"""Compute S0_ij between neighboring particles."""
self.S0 = {}
self.initial_neighbors = {}
tree = cKDTree(np.column_stack((self.X, self.Y)))
R = max(self.dx, self.dy) * self.search_radius_factor
for i in range(self.N):
neighbor_list = []
for j in tree.query_ball_point([self.X[i], self.Y[i]], R):
if j <= i:
continue
dist = np.hypot(self.X[i] - self.X[j], self.Y[i] - self.Y[j])
S0_ij = 0.5 * (self.H_initial[i] + self.H_initial[j]) * dist
self.S0[(i, j)] = S0_ij
self.S0[(j, i)] = S0_ij
neighbor_list.append(j)
self.initial_neighbors[i] = neighbor_list
# Also add reverse connections
for j in neighbor_list:
if i not in self.initial_neighbors.get(j, []):
self.initial_neighbors.setdefault(j, []).append(i)
def _compute_impact_parameter(self, xi, yi, ui, vi, xj, yj, uj, vj):
"""Return cos(theta) between particles i and j."""
dx, dy = xj - xi, yj - yi
dist = np.hypot(dx, dy) + 1e-14
du, dv = uj - ui, vj - vi
v_rel = np.hypot(du, dv) + 1e-14
return (dx * du + dy * dv) / (dist * v_rel)
def predictor_step(self):
"""Advance particle coordinates AND density by RK2 scheme."""
H, X, Y = self.H, self.X, self.Y
# Step 1: Compute velocity and divergence at t
U1, V1 = self.velocity_field(X, Y, H)
div1 = self._compute_divergence_field(X, Y, H)
# Half-step for X, Y, H
X_half = X + 0.5 * self.dt * U1
Y_half = Y + 0.5 * self.dt * V1
H_half = H * (1.0 - 0.5 * self.dt * div1) # NEW: evolve H
# Step 2: Compute velocity at half-step
U2, V2 = self.velocity_field(X_half, Y_half, H_half)
div2 = self._compute_divergence_field(X_half, Y_half, H_half)
# Full step
self.X_new = X + self.dt * U2
self.Y_new = Y + self.dt * V2
self.H_pred = H * (1.0 - self.dt * div2) # NEW: predicted H
self.U_curr, self.V_curr = U2, V2
def _compute_divergence_field(self, X, Y, H):
"""Compute analytical divergence for known velocity fields"""
if self.velocity_type == "quasilinear_rotation_pressure":
# Analytical: div(v) = 4*alpha/R²
alpha = 0.05
return np.full_like(H, 4.0 * alpha / self.R**2)
elif self.velocity_type == "rankine_vortex":
# div(v) = 0 for incompressible vortex
return np.zeros_like(H)
elif self.velocity_type == "linear_rotation":
# div(v) = 0 for solid body rotation
return np.zeros_like(H)
elif self.velocity_type == "quasilinear_transport":
# v = (0.5*H, 0.5*H)
# div(v) = 0.5*(∂H/∂x + ∂H/∂y)
# Numerical computation needed
return self._compute_divergence_numerical(X, Y, H)
else:
# Fallback to numerical
return self._compute_divergence_numerical(X, Y, H)
def _compute_divergence_numerical(self, X, Y, H):
"""Numerical divergence computation"""
div_v = np.zeros_like(H)
for i in range(self.N):
neighbors = self.initial_neighbors.get(i, [])
if len(neighbors) < 3:
continue
A = []
b_u = []
b_v = []
U, V = self.velocity_field(X, Y, H)
for j in neighbors:
dx = X[j] - X[i]
dy = Y[j] - Y[i]
du = U[j] - U[i]
dv = V[j] - V[i]
A.append([dx, dy])
b_u.append(du)
b_v.append(dv)
try:
A = np.array(A)
grad_u = np.linalg.lstsq(A, np.array(b_u), rcond=None)[0]
grad_v = np.linalg.lstsq(A, np.array(b_v), rcond=None)[0]
div_v[i] = grad_u[0] + grad_v[1]
except:
pass
return div_v
def _compute_are_effective_radius(self, i, X, Y):
neighbors = self.initial_neighbors.get(i, [])
if len(neighbors) == 0:
return None
r_avg = 0.0
for j in neighbors:
r = np.hypot(X[i] - X[j], Y[i] - Y[j])
r_avg += r
r_avg /= len(neighbors)
return np.pi * r_avg**2
def corrector_step(self):
H_new = self.H_pred.copy()
S_prev = {}
for i in range(self.N):
if i not in self.initial_neighbors:
continue
for j in self.initial_neighbors[i]:
if (i, j) in S_prev: # Already computed
continue
# Distance at OLD positions (before this step)
dist_old = np.hypot(self.X[i] - self.X[j], self.Y[i] - self.Y[j])
S_prev[(i, j)] = 0.5 * (self.H[i] + self.H[j]) * dist_old
S_prev[(j, i)] = S_prev[(i, j)]
corrections_applied = 0
all_cos_values = []
triggering_cos_values = []
passing_sign_check = 0
for i in range(self.N):
if i not in self.initial_neighbors:
continue
neighbors = self.initial_neighbors[i]
if len(neighbors) == 0:
continue
best_j = -1
min_cos = 2.0
for j in neighbors:
cos_t = self._compute_impact_parameter(
self.X_new[i],
self.Y_new[i],
self.U_curr[i],
self.V_curr[i],
self.X_new[j],
self.Y_new[j],
self.U_curr[j],
self.V_curr[j],
)
all_cos_values.append(cos_t)
if cos_t < min_cos:
min_cos, best_j = cos_t, j
if best_j == -1 or min_cos >= self.cos_theta_threshold:
continue
triggering_cos_values.append(min_cos)
S_target = S_prev[(i, best_j)]
dist_new = np.hypot(
self.X_new[i] - self.X_new[best_j], self.Y_new[i] - self.Y_new[best_j]
)
if dist_new < 1e-14:
continue
S_current = 0.5 * (self.H[i] + self.H[best_j]) * dist_new
if (self.H[i] - self.H[best_j]) * (S_current - S_target) <= 0:
continue
passing_sign_check += 1
h_corr = (2.0 * S_target / dist_new) - self.H[best_j]
h_corr = max(h_corr, 0.0)
damping = 1.0
H_new[i] = (1.0 - damping) * self.H[i] + damping * h_corr
corrections_applied += 1
if hasattr(self, "_step") and self._step % 100 == 0:
print(f"\n=== Step {self._step} Diagnostics ===")
print(f"All cos(θ) values:")
print(f" Min: {np.min(all_cos_values):.4f}")
print(f" Max: {np.max(all_cos_values):.4f}")
print(f" Mean: {np.mean(all_cos_values):.4f}")
print(f" Median: {np.median(all_cos_values):.4f}")
print(
f"\nParticles passing cos threshold (<{self.cos_theta_threshold}): {len(triggering_cos_values)}"
)
if len(triggering_cos_values) > 0:
print(
f" Their cos(θ) range: [{np.min(triggering_cos_values):.4f}, {np.max(triggering_cos_values):.4f}]"
)
print(f"Particles passing sign check: {passing_sign_check}")
print(f"Final corrections applied: {corrections_applied}\n")
if corrections_applied > 0:
print(f" Applied {corrections_applied} corrections")
self.X = self.X_new.copy()
self.Y = self.Y_new.copy()
self.H = H_new
def compute_total_mass(self):
"""Return total mass in the system."""
return np.sum(self.H) * self.dx * self.dy
def check_cfl(self):
"""Compute CFL number and velocity statistics."""
U, V = self.velocity_field(self.X, self.Y, self.H)
v_max = np.hypot(U, V).max()
h_min = min(self.dx, self.dy)
return (v_max * self.dt / h_min if v_max > 0 else 0.0), v_max, h_min
def solve(self, t_final=10.0, save_every=10, verbose=True):
"""Integrate system until t_final."""
t, results = 0.0, {
"times": [0.0],
"X": [self.X.copy()],
"Y": [self.Y.copy()],
"H": [self.H.copy()],
"masses": [self.initial_total_mass],
"cfl": [],
}
cfl, v_max, h_min = self.check_cfl()
if verbose:
print(f"Start: CFL={cfl:.4f}, v_max={v_max:.4f}")
for step in range(1, int(np.ceil(t_final / self.dt)) + 1):
self._step = step
self.predictor_step()
self.corrector_step()
t += self.dt
cfl, v_max, _ = self.check_cfl()
if step % save_every == 0 or step == int(np.ceil(t_final / self.dt)):
m = self.compute_total_mass()
results["times"].append(t)
results["X"].append(self.X.copy())
results["Y"].append(self.Y.copy())
results["H"].append(self.H.copy())
results["masses"].append(m)
results["cfl"].append(cfl)
if verbose and (
step % (save_every * 10) == 0
or step == int(np.ceil(t_final / self.dt))
):
err = (m - self.initial_total_mass) / self.initial_total_mass * 100
print(
f"Step {step:5d} | t={t:6.2f} | max(H)={self.H.max():.4f} | mass_err={err:+.25f}% | CFL={cfl:.4f}"
)
if verbose:
print(f"Done: t={t:.2f}, steps={step}")
return results
