Add files via upload

Author: Kyan0124

functions.py

@@ -0,0 +1,114 @@
+import numpy as np
+import matplotlib.pyplot as plt
+
+# ======================
+# 双线性插值采样函数
+# ======================
+def bilinear_sample(arr, yf, xf):
+    """在浮点索引 (yf=row, xf=col) 上采样 2D 数组 arr 的双线性插值值。"""
+    rows, cols = arr.shape
+    if yf < 0 or xf < 0 or yf > rows - 1 or xf > cols - 1:
+        return np.nan
+    i0 = int(np.floor(yf)); j0 = int(np.floor(xf))
+    i1 = min(i0 + 1, rows - 1); j1 = min(j0 + 1, cols - 1)
+    dy = yf - i0; dx = xf - j0
+    v00 = arr[i0, j0]; v10 = arr[i1, j0]; v01 = arr[i0, j1]; v11 = arr[i1, j1]
+    return (v00 * (1 - dy) * (1 - dx) +
+            v10 * dy * (1 - dx) +
+            v01 * (1 - dy) * dx +
+            v11 * dy * dx)
+
+
+# ======================
+# 稳定版落石物理模拟
+# ======================
+def rockfall_physics_stable(dem, start_idx, cellsize=1.0,
+                            g=9.81, dt=0.1, friction=0.3,
+                            max_steps=20000, speed_tol=1e-2, slope_tol=1e-3):
+    """
+    稳定版落石运动模拟，包含半隐式积分与低坡度停止判断。
+    返回路径 (row, col)、速度数组、动能数组。
+    """
+    rows, cols = dem.shape
+    grad_y_full, grad_x_full = np.gradient(dem, cellsize, cellsize)
+
+    start_row, start_col = start_idx
+    pos_y = start_row * cellsize
+    pos_x = start_col * cellsize
+    vx, vy = 0.0, 0.0
+
+    path_indices = [(start_row, start_col)]
+    path_speed = [0.0]
+    path_energy = [0.0]  # 动能 (0.5 * v^2)，假设单位质量=1
+
+    for step in range(max_steps):
+        frac_row = pos_y / cellsize
+        frac_col = pos_x / cellsize
+
+        if frac_row < 1 or frac_row > rows - 2 or frac_col < 1 or frac_col > cols - 2:
+            break
+
+        dg_dx = bilinear_sample(grad_x_full, frac_row, frac_col)
+        dg_dy = bilinear_sample(grad_y_full, frac_row, frac_col)
+        if np.isnan(dg_dx) or np.isnan(dg_dy):
+            break
+
+        grad_norm = np.hypot(dg_dx, dg_dy)
+        denom = np.sqrt(1.0 + grad_norm**2)
+
+        # 半隐式积分
+        a_x = -g * dg_dx / denom - friction * vx
+        a_y = -g * dg_dy / denom - friction * vy
+
+        vx += a_x * dt
+        vy += a_y * dt
+        pos_x += vx * dt
+        pos_y += vy * dt
+
+        speed = np.hypot(vx, vy)
+        kinetic_E = 0.5 * speed**2
+
+        path_indices.append((pos_y / cellsize, pos_x / cellsize))
+        path_speed.append(speed)
+        path_energy.append(kinetic_E)
+
+        if speed < speed_tol or grad_norm < slope_tol:
+            break
+
+    return np.array(path_indices), np.array(path_speed), np.array(path_energy)
+
+
+# ======================
+# 绘图函数（在独立弹窗中显示）
+# ======================
+def plot_with_speed_and_energy(dem, path_idx, speed, energy, start_idx, cmap='terrain'):
+    plt.figure(figsize=(14, 6))
+    gs = plt.GridSpec(1, 2, width_ratios=[2, 1])
+
+    # ---- 左图：DEM + 路径 ----
+    ax1 = plt.subplot(gs[0])
+    im = ax1.imshow(dem, cmap=cmap, origin='upper')
+    plt.colorbar(im, ax=ax1, label='Elevation (m)')
+
+    if len(path_idx) > 1:
+        for k in range(len(path_idx)-1):
+            y0, x0 = path_idx[k]
+            y1, x1 = path_idx[k+1]
+            s = speed[k]
+            ax1.plot([x0, x1], [y0, y1], color=plt.cm.autumn_r(min(1.0, s/10.0)), linewidth=2)
+
+    ax1.scatter([start_idx[1]], [start_idx[0]], c='k', s=60, label='Start', zorder=3)
+    ax1.set_title('Rockfall Path (color ≈ speed)')
+    ax1.legend()
+    ax1.invert_yaxis()
+
+    # ---- 右图：动能曲线 ----
+    ax2 = plt.subplot(gs[1])
+    ax2.plot(energy, color='orange', linewidth=2)
+    ax2.set_title("Kinetic Energy vs Step")
+    ax2.set_xlabel("Step")
+    ax2.set_ylabel("Energy (J/kg)")
+    ax2.grid(alpha=0.3)
+
+    plt.tight_layout()
+    plt.show()  # 原生 Matplotlib 弹窗（含放大/缩放/回退功能）

main.py

@@ -0,0 +1,101 @@
+import tkinter as tk
+from tkinter import ttk, filedialog, messagebox
+import numpy as np
+from functions import rockfall_physics_stable, plot_with_speed_and_energy
+
+class RockfallGUI:
+    def __init__(self, root):
+        self.root = root
+        self.root.title("Rockfall Simulation")
+        self.root.geometry("350x400")
+
+        # 左侧控制区
+        control_frame = ttk.Frame(root, width=300)
+        control_frame.pack(fill='both', expand=True, padx=10, pady=10)
+
+        ttk.Label(
+            control_frame,
+            text="Rockfall Simulation Parameters",
+            font=('Arial', 12, 'bold')
+        ).pack(pady=5)
+
+        # 参数变量
+        self.cellsize = tk.DoubleVar(value=30.0)
+        self.dt = tk.DoubleVar(value=0.1)
+        self.friction = tk.DoubleVar(value=0.15)
+        self.max_steps = tk.IntVar(value=5000)
+        self.start_row = tk.IntVar(value=357)
+        self.start_col = tk.IntVar(value=100)
+        self.g = tk.DoubleVar(value=9.81)
+
+        # 创建输入框
+        self.entries = {}
+        self._add_param(control_frame, "Cellsize (m):", self.cellsize)
+        self._add_param(control_frame, "Time Step (s):", self.dt)
+        self._add_param(control_frame, "Friction:", self.friction)
+        self._add_param(control_frame, "Max Steps:", self.max_steps)
+        self._add_param(control_frame, "Start Row:", self.start_row)
+        self._add_param(control_frame, "Start Col:", self.start_col)
+        self._add_param(control_frame, "Gravity:", self.g)
+
+        # 操作按钮
+        ttk.Button(control_frame, text="Load DEM", command=self.load_dem).pack(pady=10, fill='x')
+        ttk.Button(control_frame, text="Run Simulation", command=self.run_simulation).pack(pady=5, fill='x')
+
+        self.dem = None
+
+    def _add_param(self, parent, label, variable):
+        frame = ttk.Frame(parent)
+        frame.pack(fill='x', pady=3)
+        ttk.Label(frame, text=label, width=15).pack(side='left')
+        entry = ttk.Entry(frame, textvariable=variable)
+        entry.pack(side='right', fill='x', expand=True)
+        self.entries[label] = entry  # 存起来备用
+
+    def load_dem(self):
+        filepath = filedialog.askopenfilename(
+            title="Select DEM File",
+            filetypes=[("ASCII DEM (*.asc)", "*.asc"), ("All Files", "*.*")]
+        )
+        if not filepath:
+            return
+
+        try:
+            self.dem = np.loadtxt(filepath, skiprows=6)
+            messagebox.showinfo("Success", f"DEM loaded: shape = {self.dem.shape}")
+        except Exception as e:
+            messagebox.showerror("Error", f"Failed to load DEM:\n{e}")
+
+    def run_simulation(self):
+        if self.dem is None:
+            messagebox.showwarning("No DEM", "Please load a DEM file first.")
+            return
+
+        try:
+            cellsize = self.cellsize.get()
+            dt = self.dt.get()
+            friction = self.friction.get()
+            max_steps = self.max_steps.get()
+            start_row = self.start_row.get()
+            start_col = self.start_col.get()
+            g = self.g.get()
+        except tk.TclError:
+            messagebox.showerror("Invalid Input", "Please check your parameter values.")
+            return
+
+        start = (start_row, start_col)
+
+        try:
+            path_idx, speed, energy = rockfall_physics_stable(
+                self.dem, start,
+                cellsize=cellsize, g=g, dt=dt,
+                friction=friction, max_steps=max_steps
+            )
+            plot_with_speed_and_energy(self.dem, path_idx, speed, energy, start)
+        except Exception as e:
+            messagebox.showerror("Error", f"Simulation failed:\n{e}")
+
+if __name__ == "__main__":
+    root = tk.Tk()
+    app = RockfallGUI(root)
+    root.mainloop()