Revise paste rabotaet dlya evolucii

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