rishab1090/simulator
0
1# app.py2import os, uuid, zipfile, math3from pathlib import Path4from typing import Optional5from fastapi import FastAPI, UploadFile, File, Form, HTTPException6from fastapi.responses import FileResponse7import numpy as np8import rasterio9from rasterio.transform import Affine10from scipy.ndimage import gaussian_filter, map_coordinates, binary_fill_holes, binary_closing, binary_dilation, generate_binary_structure11from scipy.ndimage import label12import matplotlib.pyplot as plt13from tqdm import tqdm14from shapely.geometry import LineString, mapping15import geopandas as gpd16from fastapi.responses import RedirectResponse17 18# gravity19g = 9.8120 21app = FastAPI(title="Runout simulator single-file API")22BASE_WORKDIR = Path("/tmp/runout_jobs_single")23 24BASE_WORKDIR.mkdir(parents=True, exist_ok=True)25 26# ---------------- IO helpers ----------------27def read_dem(path):28 with rasterio.open(path) as ds:29 dem = ds.read(1).astype(float)30 meta = ds.meta.copy()31 transform = ds.transform32 crs = ds.crs33 return dem, transform, crs, meta34 35def save_geotiff(arr, meta, path):36 meta2 = meta.copy()37 meta2.update(dtype="float32", count=1)38 os.makedirs(os.path.dirname(path) or ".", exist_ok=True)39 with rasterio.open(path, "w", **meta2) as dst:40 dst.write(np.array(arr, dtype='float32'), 1)41 42def save_png(arr, out_path, cmap="inferno", smooth=1.0, log_scale=False, clip_percent=(0.1,99.9)):43 a = np.array(arr, dtype=float)44 a[~np.isfinite(a)] = np.nan45 if smooth and smooth > 0:46 a = gaussian_filter(a, sigma=smooth)47 vals = a[~np.isnan(a)]48 if vals.size == 0:49 a = np.zeros((10,10))50 vmin, vmax = 0,151 else:52 if log_scale:53 a = np.log1p(np.clip(a,0,None)*1e4)54 vals = a[~np.isnan(a)]55 vmin = float(np.percentile(vals, clip_percent[0]))56 vmax = float(np.percentile(vals, clip_percent[1]))57 if vmax <= vmin:58 vmax = vmin + 1e-659 os.makedirs(os.path.dirname(out_path) or ".", exist_ok=True)60 plt.imsave(out_path, a, cmap=cmap, vmin=vmin, vmax=vmax)61 62def save_samples_geojson(samples, transform, out_path, crs=None):63 feats=[]64 for i,path in enumerate(samples):65 if not path or len(path)<2: continue66 coords=[rasterio.transform.xy(transform,int(r),int(c)) for (r,c) in path]67 feats.append({"type":"Feature","geometry":mapping(LineString(coords)),"properties":{"id":i,"steps":len(coords)}})68 if not feats:69 return70 gdf = gpd.GeoDataFrame.from_features(feats)71 if crs is not None:72 try:73 gdf.set_crs(crs, inplace=True)74 except Exception:75 pass76 os.makedirs(os.path.dirname(out_path) or ".", exist_ok=True)77 gdf.to_file(out_path, driver="GeoJSON")78 79# ---------------- interpolation & gradients ----------------80def interp_bilinear(arr, pts):81 res = map_coordinates(arr, pts, order=1, mode='nearest')82 if np.size(res) == 1:83 return float(res)84 return res85 86def dem_gradients(dem, cellsize):87 # return dz/dx, dz/dy (rise per meter)88 dy, dx = np.gradient(dem, cellsize, cellsize)89 return dx, dy90 91def compute_slope_field(dx_field, dy_field):92 grad_mag = np.sqrt(dx_field*dx_field + dy_field*dy_field)93 slope_deg = np.degrees(np.arctan(grad_mag))94 return slope_deg95 96# ---------------- rim detection ----------------97def detect_rim_coords(dem, cellsize, pit_depth_frac=0.25, rim_buffer_m=6.0, min_rim_samples=500):98 dem_valid = np.where(np.isfinite(dem), dem, np.nan)99 zmin = float(np.nanmin(dem_valid)); zmax = float(np.nanmax(dem_valid))100 depth = zmax - zmin101 if depth <= 0:102 h,w = dem.shape103 return [(h//2,w//2)], depth104 pit_thresh = zmin + pit_depth_frac * depth105 pit_mask = (dem <= pit_thresh) & np.isfinite(dem)106 pit_mask = binary_fill_holes(pit_mask)107 pit_mask = binary_closing(pit_mask, structure=generate_binary_structure(2,2), iterations=2)108 lbls, nlab = label(pit_mask)109 if nlab > 1:110 counts = [(lbls==i).sum() for i in range(1, nlab+1)]111 largest = int(np.argmax(counts) + 1)112 pit_mask = (lbls == largest)113 iters = max(1, int(round(rim_buffer_m / cellsize)))114 rim_band = binary_dilation(pit_mask, iterations=iters) & (~pit_mask)115 coords = list(zip(*np.where(rim_band)))116 if len(coords) > min_rim_samples:117 step = max(1, int(len(coords)/min_rim_samples))118 coords = coords[::step]119 return coords, depth120 121# ---------------- steepest march ----------------122def march_steepest(dem, start_rc, cellsize, slope_deg, slope_thresh_deg, min_drop_m, max_steps):123 offs = [(-1,-1),(-1,0),(-1,1),(0,-1),(0,1),(1,-1),(1,0),(1,1)]124 rows, cols = dem.shape125 r, c = int(start_rc[0]), int(start_rc[1])126 z0 = dem[r,c]127 visited = set(); visited.add((r,c))128 for step in range(max_steps):129 best = None; best_dz = 0.0130 z = dem[r,c]131 if not np.isfinite(z):132 break133 for (dr,dc) in offs:134 rr,cc = r+dr, c+dc135 if rr<0 or rr>=rows or cc<0 or cc>=cols: continue136 if (rr,cc) in visited: continue137 zn = dem[rr,cc]138 if not np.isfinite(zn): continue139 dz = z - zn140 if dz > best_dz:141 best_dz = dz; best = (rr,cc)142 if best is None or best_dz <= 0:143 break144 r,c = best; visited.add((r,c))145 cum_drop = z0 - dem[r,c]146 if slope_deg[r,c] >= slope_thresh_deg and cum_drop >= min_drop_m:147 return (r,c), cum_drop, step+1148 return None, 0.0, step+1149 150# ---------------- particle integrator with wall handling ----------------151def particle_run(dem, dx_field, dy_field, cellsize, start_xy_m, initial_speed, mu, step_loss,152 transform, xi=1000.0, wall_slope_thresh_deg=60.0, wall_restitution=0.2,153 wall_deflect_frac=0.85, stop_on_wall_ke=5.0,154 dt_max=0.2, dt_min=0.005, max_steps=20000, v_threshold=0.2):155 x_m, y_m = start_xy_m156 rows, cols = dem.shape157 inv = ~transform158 159 # initial gradient / direction160 col0, row0 = inv * (x_m, y_m)161 pts0 = np.array([[row0],[col0]])162 try:163 gx0 = float(interp_bilinear(dx_field, pts0))164 gy0 = float(interp_bilinear(dy_field, pts0))165 except Exception:166 gx0 = 0.0; gy0 = 0.0167 downslope = np.array([-gx0, -gy0]); n0 = np.linalg.norm(downslope)168 dir_x, dir_y = (downslope / n0) if n0 > 1e-9 else (0.0, 0.0)169 170 vx = initial_speed * dir_x171 vy = initial_speed * dir_y172 pos_x = x_m; pos_y = y_m173 174 visited = []175 ke_cells = {}176 177 for step in range(max_steps):178 col_cur, row_cur = inv * (pos_x, pos_y)179 if col_cur < 0 or col_cur >= cols-1 or row_cur < 0 or row_cur >= rows-1:180 break181 182 pts = np.array([[row_cur],[col_cur]])183 try:184 gz_x = float(interp_bilinear(dx_field, pts))185 gz_y = float(interp_bilinear(dy_field, pts))186 except Exception:187 gz_x = 0.0; gz_y = 0.0188 189 grad_mag = math.hypot(gz_x, gz_y)190 local_slope_deg = math.degrees(math.atan(grad_mag))191 192 a_drive_x = -g * gz_x193 a_drive_y = -g * gz_y194 195 vmag = math.hypot(vx, vy)196 if vmag > 1e-9:197 ux = vx / vmag; uy = vy / vmag198 a_voellmy_mag = mu * g + (vmag * vmag) / max(1e-12, xi)199 a_voellmy_x = -a_voellmy_mag * ux200 a_voellmy_y = -a_voellmy_mag * uy201 else:202 a_voellmy_x = a_voellmy_y = 0.0203 204 c_drag = 0.03205 a_drag_x = -c_drag * vx * vmag206 a_drag_y = -c_drag * vy * vmag207 208 ax = a_drive_x + a_voellmy_x + a_drag_x209 ay = a_drive_y + a_voellmy_y + a_drag_y210 211 dt = min(dt_max, max(dt_min, 0.4 * (cellsize / (vmag + 1e-6))))212 213 vx_mid = vx + 0.5 * ax * dt214 vy_mid = vy + 0.5 * ay * dt215 new_pos_x = pos_x + vx_mid * dt216 new_pos_y = pos_y + vy_mid * dt217 new_vx = vx + ax * dt218 new_vy = vy + ay * dt219 220 # wall check at next location221 col_next, row_next = inv * (new_pos_x, new_pos_y)222 wall_encounter = False223 gz_x_n = gz_y_n = grad_mag_n = 0.0224 if 0 <= col_next < cols and 0 <= row_next < rows:225 try:226 gz_x_n = float(interp_bilinear(dx_field, np.array([[row_next],[col_next]])))227 gz_y_n = float(interp_bilinear(dy_field, np.array([[row_next],[col_next]])))228 grad_mag_n = math.hypot(gz_x_n, gz_y_n)229 local_slope_deg_n = math.degrees(math.atan(grad_mag_n))230 except Exception:231 local_slope_deg_n = 0.0232 grad_mag_n = 0.0233 if local_slope_deg_n >= wall_slope_thresh_deg:234 wall_encounter = True235 n_hat = np.array([gz_x_n, gz_y_n]) / max(1e-12, grad_mag_n)236 237 if wall_encounter:238 ke = 0.5 * (vmag * vmag)239 if ke <= stop_on_wall_ke:240 r_idx = int(round(row_cur)); c_idx = int(round(col_cur))241 if 0 <= r_idx < rows and 0 <= c_idx < cols:242 visited.append((r_idx, c_idx))243 ke_cells[(r_idx,c_idx)] = max(ke_cells.get((r_idx,c_idx), 0.0), ke)244 break245 246 t = np.array([-gz_y_n, gz_x_n])247 tnorm = np.linalg.norm(t)248 if tnorm < 1e-12:249 v_vec = np.array([new_vx, new_vy])250 nv = np.dot(v_vec, n_hat)251 v_reflect = v_vec - (1.0 + wall_restitution) * nv * n_hat252 v_reflect *= 0.9253 new_vx, new_vy = float(v_reflect[0]), float(v_reflect[1])254 else:255 t_hat = t / tnorm256 v_vec = np.array([new_vx, new_vy])257 v_tang = np.dot(v_vec, t_hat) * t_hat258 v_norm = v_vec - v_tang259 new_v_vec = wall_deflect_frac * v_tang - wall_restitution * v_norm260 new_vx, new_vy = float(new_v_vec[0]), float(new_v_vec[1])261 new_pos_x = pos_x + new_vx * dt * 0.6262 new_pos_y = pos_y + new_vy * dt * 0.6263 264 vmag_after = math.hypot(new_vx, new_vy)265 if vmag_after < v_threshold:266 r_idx = int(round(row_cur)); c_idx = int(round(col_cur))267 if 0 <= r_idx < rows and 0 <= c_idx < cols:268 visited.append((r_idx, c_idx))269 ke_cells[(r_idx,c_idx)] = max(ke_cells.get((r_idx,c_idx), 0.0), 0.5 * vmag_after * vmag_after)270 break271 272 pos_x, pos_y = new_pos_x, new_pos_y273 vx, vy = new_vx, new_vy274 else:275 pos_x, pos_y = new_pos_x, new_pos_y276 vx, vy = new_vx, new_vy277 278 loss = np.clip(np.random.normal(loc=step_loss, scale=step_loss*0.2), 0.0, 0.6)279 vx *= (1.0 - loss); vy *= (1.0 - loss)280 vmag = math.hypot(vx, vy)281 282 col_round, row_round = inv * (pos_x, pos_y)283 r_idx = int(round(row_round)); c_idx = int(round(col_round))284 if r_idx < 0 or r_idx >= rows or c_idx < 0 or c_idx >= cols:285 break286 visited.append((r_idx, c_idx))287 ke_local = 0.5 * (vmag * vmag)288 ke_cells[(r_idx, c_idx)] = max(ke_cells.get((r_idx, c_idx), 0.0), ke_local)289 290 if vmag < v_threshold:291 break292 293 unique_visited = []294 seen = set()295 for rc in visited:296 if rc not in seen:297 unique_visited.append(rc); seen.add(rc)298 return unique_visited, ke_cells299 300# ---------------- Monte Carlo ----------------301def run_mc(dem, transform, release_cells, cellsize, trials, mu_mean, mu_std, xi, step_loss,302 min_drop, force_drop, sample_paths, seed, wall_slope_thresh_deg, wall_restitution, wall_deflect_frac, stop_on_wall_ke):303 rng = np.random.default_rng(seed)304 dx_field, dy_field = dem_gradients(dem, cellsize)305 heat = np.zeros_like(dem, dtype=float)306 impact = np.zeros_like(dem, dtype=float)307 sample_paths_out = []308 n_releases = len(release_cells)309 if n_releases == 0:310 raise RuntimeError("No release cells provided.")311 for t in tqdm(range(trials), desc="Trials"):312 r,c = release_cells[rng.integers(0, n_releases)]313 x0, y0 = rasterio.transform.xy(transform, int(r), int(c))314 if force_drop is not None:315 drop_m = float(force_drop)316 else:317 drop_m = max(min_drop, 0.02 * (np.nanmax(dem) - np.nanmin(dem)))318 initial_speed = math.sqrt(2.0 * g * drop_m)319 mu = float(max(0.0, rng.normal(mu_mean, mu_std)))320 path, ke_map = particle_run(dem, dx_field, dy_field, cellsize,321 (x0, y0), initial_speed, mu,322 step_loss, transform,323 xi=xi,324 wall_slope_thresh_deg=wall_slope_thresh_deg,325 wall_restitution=wall_restitution,326 wall_deflect_frac=wall_deflect_frac,327 stop_on_wall_ke=stop_on_wall_ke)328 for (rr, cc) in path:329 heat[rr, cc] += 1.0330 if (rr,cc) in ke_map:331 impact[rr, cc] = max(impact[rr, cc], ke_map[(rr,cc)])332 if len(sample_paths_out) < sample_paths:333 sample_paths_out.append(path)334 heat /= float(max(1, trials))335 return heat, impact, sample_paths_out336 337# ---------------- utils ----------------338def print_stats(name, arr):339 a = np.array(arr, dtype=float)340 a[~np.isfinite(a)] = np.nan341 non = a[~np.isnan(a)]342 if non.size == 0:343 return f"[STAT] {name}: no finite values"344 return f"[STAT] {name}: min={np.nanmin(non):.6g}, p1={np.nanpercentile(non,1):.6g}, p50={np.nanpercentile(non,50):.6g}, p99={np.nanpercentile(non,99):.6g}, max={np.nanmax(non):.6g}, mean={np.nanmean(non):.6g}"345 346# ---------------- API endpoint ----------------347@app.post("/run")348async def run_endpoint(349 dem: UploadFile = File(...),350 trials: int = Form(500),351 release_density: int = Form(1),352 pit_depth_frac_release: float = Form(0.25),353 slope_thresh_deg: float = Form(5.0),354 mu_mean: float = Form(0.05),355 mu_std: float = Form(0.01),356 xi: float = Form(1000.0),357 step_loss: float = Form(0.001),358 sample_paths: int = Form(50),359 wall_slope_thresh_deg: float = Form(60.0),360 wall_restitution: float = Form(0.2),361 wall_deflect_frac: float = Form(0.85),362 stop_on_wall_ke: float = Form(5.0),363 seed: int = Form(0),364 out_prefix: Optional[str] = Form(None),365):366 job_id = uuid.uuid4().hex[:12]367 job_dir = BASE_WORKDIR / job_id368 job_dir.mkdir(parents=True, exist_ok=True)369 dem_path = job_dir / "input_dem.tif"370 with dem_path.open("wb") as f:371 f.write(await dem.read())372 # read dem373 try:374 dem_arr, transform, crs, meta = read_dem(str(dem_path))375 except Exception as e:376 raise HTTPException(status_code=400, detail=f"Failed reading DEM: {e}")377 cellsize = abs(transform.a)378 # detect rim379 rim_coords, pit_depth = detect_rim_coords(dem_arr, cellsize, pit_depth_frac=pit_depth_frac_release, rim_buffer_m=6.0)380 dx_field, dy_field = dem_gradients(dem_arr, cellsize)381 slope_field = compute_slope_field(dx_field, dy_field)382 # march steepest to find release targets383 release_cells=[]384 for rc in rim_coords:385 tgt, cum_drop, steps = march_steepest(dem_arr, rc, cellsize, slope_field, slope_thresh_deg, 2.0, 400)386 if tgt is not None:387 release_cells.append(tgt)388 seeds = release_cells * max(1, int(release_density))389 if len(seeds) == 0:390 raise HTTPException(status_code=400, detail="No release seeds found.")391 # debug rim png392 debug_rim_png = job_dir / f"{job_id}_rim_debug.png"393 try:394 overlay = np.zeros((dem_arr.shape[0], dem_arr.shape[1], 3), dtype=float)395 dmmin = float(np.nanmin(dem_arr)); dmmax = float(np.nanmax(dem_arr))396 norm = (dem_arr - dmmin) / (dmmax - dmmin + 1e-9)397 overlay[..., :] = np.expand_dims(norm,2) * 0.35398 for (r,c) in release_cells:399 if 0 <= r < overlay.shape[0] and 0 <= c < overlay.shape[1]:400 overlay[r,c] = [1.0, 0.0, 0.0]401 plt.imsave(str(debug_rim_png), np.clip(overlay,0,1))402 except Exception:403 pass404 # run Monte Carlo405 try:406 heat, impact, sample_paths = run_mc(dem_arr, transform, seeds, cellsize,407 int(trials), float(mu_mean), float(mu_std), float(xi), float(step_loss),408 float(2.0), None, int(sample_paths), int(seed),409 float(wall_slope_thresh_deg), float(wall_restitution), float(wall_deflect_frac), float(stop_on_wall_ke))410 except Exception as e:411 raise HTTPException(status_code=500, detail=f"Simulation error: {e}")412 # outputs413 prefix = out_prefix or job_id414 out_heat_png = job_dir / f"{prefix}_heat.png"415 out_impact_png = job_dir / f"{prefix}_impact.png"416 out_heat_tif = job_dir / f"{prefix}_heat.tif"417 out_impact_tif = job_dir / f"{prefix}_impact.tif"418 out_samples_geojson = job_dir / f"{prefix}_samples.geojson"419 try:420 save_geotiff(heat, meta, str(out_heat_tif))421 save_geotiff(impact, meta, str(out_impact_tif))422 save_png(heat, str(out_heat_png), cmap="inferno", smooth=1.0, log_scale=False)423 save_png(impact, str(out_impact_png), cmap="magma", smooth=1.0, log_scale=True)424 save_samples_geojson(sample_paths, transform, str(out_samples_geojson), crs=crs)425 except Exception as e:426 raise HTTPException(status_code=500, detail=f"Failed to write outputs: {e}")427 # diagnostics text428 diag_txt = job_dir / "diagnostics.txt"429 with diag_txt.open("w") as f:430 f.write(f"Rim pixels: {len(rim_coords)} pit_depth: {pit_depth}\n")431 f.write(print_stats("heat", heat) + "\n")432 f.write(print_stats("impact", impact) + "\n")433 try:434 idx = np.nanargmax(impact)435 r = int(idx) // impact.shape[1]; c = int(idx) % impact.shape[1]436 x,y = rasterio.transform.xy(transform, r, c)437 f.write(f"Max impact at row {r} col {c} -> x={x:.3f}, y={y:.3f}, KE={impact[r,c]:.6g}\n")438 except Exception:439 pass440 # zip outputs441 zip_path = job_dir / f"{job_id}_outputs.zip"442 with zipfile.ZipFile(zip_path, "w", compression=zipfile.ZIP_DEFLATED) as zf:443 for p in [out_heat_png, out_impact_png, out_heat_tif, out_impact_tif, out_samples_geojson, debug_rim_png, diag_txt]:444 if p.exists():445 zf.write(p, arcname=p.name)446 return FileResponse(path=str(zip_path), filename=zip_path.name, media_type="application/zip")447 448@app.get("/health")449def health():450 return {"status":"ok"}451 452 453@app.get("/")454def root():455 return RedirectResponse(url="/docs")456 457