API tools faq
paste
kosbrown

weather sim in python

kosbrown
Oct 22nd, 2025
102
0
Never
Add comment
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 43.08 KB | Gaming | 0 0
raw download clone embed print report
  1. #!/usr/bin/env python3
  2. """
  3. All-in-one Visual Weather Engine (Phases 1–3) with terminal progress bar and interactive phase picker
  4.  
  5. Phases:
  6.  1 - Static visualization: temperature heatmap + isotherms, rainfall, wind
  7.  2 - Animated particles: wind streaks and rain droplets (MP4/GIF or interactive)
  8.  3 - 3D scene: terrain surface with weather overlays (Matplotlib 3D or Plotly)
  9.  
  10. Key features:
  11.  - Ingests heightmap (grayscale or RGB) and texture (RGB).
  12.  - Prompts for sea level threshold (normalized 0..1). Default: 0.22.
  13.  - Interactive phase picker: [1/2/3/all]
  14.  - Cross-references oceans from height and texture.
  15.  - Simulates temperature (land vs ocean), divergence-free wind, humidity, rainfall.
  16.  - Timestamped outputs to avoid overwriting (shared stamp for "all").
  17.  
  18. Dependencies:
  19.  pip install numpy matplotlib pillow scipy
  20.  Optional: plotly (for interactive 3D), kaleido (for static plotly export), ffmpeg on system (for MP4 export)
  21. """
  22.  
  23. import argparse
  24. import os
  25. import sys
  26. import shutil
  27. import time
  28. import datetime as _dt
  29. import numpy as np
  30. import matplotlib as mpl
  31. import matplotlib.pyplot as plt
  32. from matplotlib import animation
  33. from matplotlib import patheffects as pe
  34. from matplotlib.collections import LineCollection
  35. from matplotlib.colors import ListedColormap, Normalize
  36. from PIL import Image
  37. from scipy.ndimage import (
  38.     gaussian_filter,
  39.     distance_transform_edt,
  40.     maximum_filter,
  41.     binary_opening,
  42.     binary_closing,
  43. )
  44. # 3D (registers projection)
  45. from mpl_toolkits.mplot3d import Axes3D  # noqa: F401
  46.  
  47.  
  48. # -----------------------------
  49. # Terminal progress bar
  50. # -----------------------------
  51.  
  52. class ProgressBar:
  53.     def __init__(self, width=40):
  54.         self.width = width
  55.         self.total = 100
  56.         self.current = 0
  57.         self.desc = ""
  58.         self._last_len = 0
  59.         self._start_time = None
  60.  
  61.     def _term_width(self):
  62.         try:
  63.             cols = shutil.get_terminal_size((80, 20)).columns
  64.             return max(40, cols)
  65.         except Exception:
  66.             return 80
  67.  
  68.     def new_task(self, desc, total):
  69.         self.desc = str(desc)
  70.         self.total = max(1, int(total))
  71.         self.current = 0
  72.         self._start_time = time.time()
  73.         self._render()
  74.  
  75.     def update(self, n):
  76.         self.current = max(0, min(int(n), int(self.total)))
  77.         self._render()
  78.  
  79.     def advance(self, dn):
  80.         self.update(self.current + dn)
  81.  
  82.     def done(self, suffix=""):
  83.         self.current = self.total
  84.         self._render(done=True, suffix=suffix)
  85.         sys.stdout.write("\n")
  86.         sys.stdout.flush()
  87.         self.desc = ""
  88.         self._start_time = None
  89.  
  90.     def _render(self, done=False, suffix=""):
  91.         pct = (float(self.current) / float(self.total)) if self.total else 1.0
  92.         pct = max(0.0, min(1.0, pct))
  93.  
  94.         term_w = self._term_width()
  95.         bar_w = min(self.width, max(10, term_w - 30))
  96.         filled = int(bar_w * pct)
  97.         bar = "=" * filled + (">" if filled < bar_w and not done else "") + "-" * max(0, bar_w - filled - (0 if filled == bar_w or done else 1))
  98.         percent_txt = f"{int(pct * 100):3d}%"
  99.         elapsed = ""
  100.         if self._start_time is not None:
  101.             dt = time.time() - self._start_time
  102.             if dt >= 0.1:
  103.                 elapsed = f"  {dt:5.1f}s"
  104.  
  105.         line = f"\r[{bar}] {percent_txt}  {self.desc}{elapsed}"
  106.         if suffix:
  107.             line += f"  {suffix}"
  108.  
  109.         pad = max(0, self._last_len - len(line))
  110.         sys.stdout.write(line + (" " * pad))
  111.         sys.stdout.flush()
  112.         self._last_len = len(line)
  113.  
  114.  
  115. # -----------------------------
  116. # Prompts / I/O helpers
  117. # -----------------------------
  118.  
  119. def prompt_for_existing_file(prompt_msg, default=None):
  120.     while True:
  121.         prompt = f"{prompt_msg}"
  122.         if default:
  123.             prompt += f" [{default}]"
  124.         prompt += ": "
  125.         path = input(prompt).strip().strip('"').strip("'")
  126.         if path == "" and default:
  127.             path = default
  128.         if os.path.isfile(path):
  129.             return path
  130.         print("File not found. Please provide a valid file path.")
  131.  
  132. def prompt_for_float(prompt_msg, default=None, minv=None, maxv=None):
  133.     while True:
  134.         prompt = f"{prompt_msg}"
  135.         if default is not None:
  136.             prompt += f" [{default}]"
  137.         prompt += ": "
  138.         s = input(prompt).strip()
  139.         if s == "" and default is not None:
  140.             val = float(default)
  141.         else:
  142.             try:
  143.                 val = float(s)
  144.             except Exception:
  145.                 print("Please enter a number.")
  146.                 continue
  147.         if minv is not None and val < minv:
  148.             print(f"Value must be >= {minv}")
  149.             continue
  150.         if maxv is not None and val > maxv:
  151.             print(f"Value must be <= {maxv}")
  152.             continue
  153.         return val
  154.  
  155. def prompt_for_phase(default_choice=None):
  156.     """
  157.    Interactive picker for phase(s). Returns a list of phases to run.
  158.    """
  159.     while True:
  160.         prompt = "Pick a phase [1/2/3/all]"
  161.         if default_choice is not None:
  162.             prompt += f" [{default_choice}]"
  163.         prompt += ": "
  164.         s = input(prompt).strip().lower()
  165.         if s == "" and default_choice is not None:
  166.             s = str(default_choice).lower()
  167.         if s in ("1", "2", "3"):
  168.             return [int(s)]
  169.         if s in ("all", "a", "123"):
  170.             return [1, 2, 3]
  171.         print("Please enter 1, 2, 3, or 'all'.")
  172.  
  173. def read_image_as_array(path, mode=None):
  174.     img = Image.open(path)
  175.     if mode:
  176.         img = img.convert(mode)
  177.     arr = np.array(img)
  178.     return arr
  179.  
  180. def ensure_gray2d(a):
  181.     """
  182.    Ensure a is a single-channel 2D array (float32).
  183.    If RGB, convert to luminance (BT.709).
  184.    """
  185.     a = np.asarray(a)
  186.     if a.ndim == 2:
  187.         return a.astype(np.float32)
  188.     if a.ndim == 3:
  189.         if a.shape[2] == 1:
  190.             return a[..., 0].astype(np.float32)
  191.         a = a.astype(np.float32)
  192.         r, g, b = a[..., 0], a[..., 1], a[..., 2]
  193.         gray = 0.2126 * r + 0.7152 * g + 0.0722 * b
  194.         return gray.astype(np.float32)
  195.     return np.squeeze(a).astype(np.float32)
  196.  
  197. def to_float01(arr):
  198.     arr = np.asarray(arr)
  199.     if arr.ndim == 2:
  200.         arr = arr.astype(np.float32)
  201.         mn, mx = float(arr.min()), float(arr.max())
  202.         if mx > mn:
  203.             arr = (arr - mn) / (mx - mn)
  204.         else:
  205.             arr = np.zeros_like(arr, dtype=np.float32)
  206.         return arr
  207.     else:
  208.         if arr.dtype == np.uint8:
  209.             return arr.astype(np.float32) / 255.0
  210.         arr = arr.astype(np.float32)
  211.         mn, mx = float(arr.min()), float(arr.max())
  212.         return (arr - mn) / (mx - mn + 1e-9)
  213.  
  214. def resize_to_match(arr, target_shape):
  215.     """
  216.    Resize arr (H,W) or (H,W,3) to target_shape (H,W[,C]) using bilinear.
  217.    Only H and W from target_shape are used.
  218.    """
  219.     if isinstance(target_shape, tuple) and len(target_shape) >= 2:
  220.         H, W = int(target_shape[0]), int(target_shape[1])
  221.     else:
  222.         raise ValueError("target_shape must be a tuple like (H, W) or (H, W, C)")
  223.  
  224.     arr_clipped = np.clip(arr, 0, 1)
  225.     im = Image.fromarray((arr_clipped * 255).astype(np.uint8))
  226.     im = im.resize((W, H), resample=Image.BILINEAR)
  227.     return (np.array(im).astype(np.float32) / 255.0)
  228.  
  229. def add_timestamp_to_path(path, fmt="%Y%m%d_%H%M%S", stamp=None):
  230.     """
  231.    Insert a timestamp before the file extension to avoid overwriting.
  232.    If stamp is provided, reuse it across multiple outputs.
  233.    """
  234.     directory, filename = os.path.split(path)
  235.     name, ext = os.path.splitext(filename)
  236.     if stamp is None:
  237.         stamp = _dt.datetime.now().strftime(fmt)
  238.     new_name = f"{name}_{stamp}{ext}"
  239.     return os.path.join(directory if directory else ".", new_name)
  240.  
  241.  
  242. # -----------------------------
  243. # Math helpers and grids
  244. # -----------------------------
  245.  
  246. def fbm_noise(shape, octaves=5, base_sigma=40.0, persistence=0.5, lacunarity=2.0, seed=42):
  247.     rng = np.random.default_rng(seed)
  248.     H, W = shape
  249.     acc = np.zeros((H, W), dtype=np.float32)
  250.     amp = 1.0
  251.     sigma = base_sigma
  252.     total_amp = 0.0
  253.     for _ in range(octaves):
  254.         n = rng.standard_normal((H, W)).astype(np.float32)
  255.         s = max(0.5, sigma)
  256.         layer = gaussian_filter(n, s, mode='reflect')
  257.         acc += amp * layer
  258.         total_amp += amp
  259.         amp *= persistence
  260.         sigma /= lacunarity
  261.     acc /= (total_amp + 1e-8)
  262.     acc = (acc - acc.min()) / (acc.max() - acc.min() + 1e-9)
  263.     return acc
  264.  
  265. def logistic_stable(x):
  266.     """
  267.    Numerically stable logistic sigmoid: sigma(x) without overflow.
  268.    """
  269.     x = x.astype(np.float64)
  270.     out = np.empty_like(x)
  271.     pos = x >= 0
  272.     with np.errstate(over='ignore', under='ignore'):
  273.         out[pos] = 1.0 / (1.0 + np.exp(-x[pos]))
  274.         ex = np.exp(x[~pos])
  275.         out[~pos] = ex / (1.0 + ex)
  276.     return out.astype(np.float32)
  277.  
  278. def compute_grids(H, W):
  279.     y = np.linspace(0.0, 1.0, H, dtype=np.float32)
  280.     x = np.linspace(0.0, 1.0, W, dtype=np.float32)
  281.     Y, X = np.meshgrid(y, x, indexing='ij')
  282.     return X, Y
  283.  
  284.  
  285. # -----------------------------
  286. # Meteorological fields
  287. # -----------------------------
  288.  
  289. def detect_water_mask_combined(texture_rgb, height_norm, sea_level=0.22,
  290.                                blue_thresh=0.45, blue_margin=0.05,
  291.                                smooth_sigma=1.2, morph=True):
  292.     """
  293.    Combine height-based and color-based water detection.
  294.      - by_height: height <= sea_level
  295.      - by_color: blue dominant over R/G by margin, above threshold
  296.    """
  297.     by_height = height_norm <= float(sea_level)
  298.  
  299.     red = texture_rgb[..., 0]
  300.     green = texture_rgb[..., 1]
  301.     blue = texture_rgb[..., 2]
  302.     by_color = (blue > blue_thresh) & (blue > green + blue_margin) & (blue > red + blue_margin)
  303.  
  304.     water = by_height | by_color
  305.     water = gaussian_filter(water.astype(np.float32), sigma=smooth_sigma) > 0.5
  306.     if morph:
  307.         water = binary_opening(water, structure=np.ones((3, 3)))
  308.         water = binary_closing(water, structure=np.ones((3, 3)))
  309.  
  310.     # Signed distance to coast: negative over water, positive over land
  311.     land = ~water
  312.     dist_to_land = distance_transform_edt(~water)     # distance for water pixels to nearest land
  313.     dist_to_water = distance_transform_edt(~land)     # distance for land pixels to nearest water
  314.     signed_coast_dist = dist_to_land.astype(np.float32)
  315.     signed_coast_dist[water] = -dist_to_water[water].astype(np.float32)
  316.  
  317.     return water, signed_coast_dist
  318.  
  319. def compute_temperature(height_norm, X, Y, water_mask, coast_dist,
  320.                         seed=42, max_elev_m=3000.0, sea_level_temp_C=22.0,
  321.                         lapse_rate_K_per_km=6.5, lat_grad_C=10.0, sst_variance_C=1.5):
  322.     """
  323.    Temperature (°C), blending land lapse-rate cooling and sea surface temperature.
  324.    """
  325.     H, W = height_norm.shape
  326.     elev_m = height_norm * max_elev_m
  327.     lapse = lapse_rate_K_per_km * (elev_m / 1000.0)
  328.  
  329.     lat_component_land = lat_grad_C * (Y - 0.5)
  330.     lat_component_sea = 0.6 * lat_grad_C * (Y - 0.5)
  331.  
  332.     land_noise = fbm_noise((H, W), octaves=5, base_sigma=min(H, W)/16,
  333.                            persistence=0.55, lacunarity=2.1, seed=seed+101)
  334.     land_noise = (land_noise - 0.5) * 4.0  # +/- 2C
  335.  
  336.     sst_noise = fbm_noise((H, W), octaves=4, base_sigma=min(H, W)/14,
  337.                           persistence=0.6, lacunarity=2.0, seed=seed+102)
  338.     sst_noise = (sst_noise - 0.5) * (2 * sst_variance_C)
  339.  
  340.     T_land = sea_level_temp_C - lapse + lat_component_land + land_noise
  341.     T_sea = sea_level_temp_C + lat_component_sea + sst_noise
  342.  
  343.     # Stable coastline blend (0 over water, 1 over land)
  344.     k = 5.0  # pixels
  345.     blend = logistic_stable(coast_dist / k)
  346.     T = blend * T_land + (1.0 - blend) * T_sea
  347.     return T
  348.  
  349. def compute_streamfunction(height_norm, X, Y, water_mask,
  350.                            base_wind_dir_deg=35.0, base_wind_speed=6.0,
  351.                            seed=42, k_orog=0.9, k_noise=1.0,
  352.                            ocean_speed_boost=0.25, mountain_drag=0.12):
  353.     """
  354.    Divergence-free flow via streamfunction with speed scaling over ocean/mountains.
  355.    """
  356.     H, W = height_norm.shape
  357.     theta = np.deg2rad(base_wind_dir_deg)
  358.     u0 = base_wind_speed * np.cos(theta)
  359.     v0 = base_wind_speed * np.sin(theta)
  360.  
  361.     elev_smooth = gaussian_filter(height_norm, sigma=min(H, W)/80)
  362.     noise = fbm_noise((H, W), octaves=5, base_sigma=min(H, W)/20, seed=seed+202)
  363.     noise = gaussian_filter(noise, sigma=min(H, W)/60)
  364.  
  365.     psi = (u0 * Y) - (v0 * X) + k_orog * elev_smooth + k_noise * noise
  366.  
  367.     dy = 1.0 / max(H - 1, 1)
  368.     dx = 1.0 / max(W - 1, 1)
  369.     dpsidy = np.gradient(psi, dy, axis=0)
  370.     dpsidx = np.gradient(psi, dx, axis=1)
  371.  
  372.     u = gaussian_filter(dpsidy, sigma=1.0)
  373.     v = gaussian_filter(-dpsidx, sigma=1.0)
  374.  
  375.     speed_scale = 1.0 + ocean_speed_boost * (water_mask.astype(np.float32)) - mountain_drag * (height_norm ** 1.5)
  376.     u *= speed_scale
  377.     v *= speed_scale
  378.  
  379.     return u, v
  380.  
  381. def compute_humidity(height_norm, water_mask, coast_dist, seed=42):
  382.     """
  383.    Humidity 0..1; high over ocean, decays inland, penalized by elevation.
  384.    """
  385.     H, W = height_norm.shape
  386.     dist = np.abs(coast_dist)
  387.     dist_norm = dist / (dist.max() + 1e-9)
  388.  
  389.     base_ocean = 0.9
  390.     base_land = 0.35
  391.     baseline = np.where(water_mask, base_ocean, base_land)
  392.  
  393.     inland_decay = np.exp(-3.0 * dist_norm)
  394.     elev_penalty = np.clip(1.0 - height_norm, 0, 1)
  395.  
  396.     noise = fbm_noise((H, W), octaves=4, base_sigma=min(H, W)/12, seed=seed+303)
  397.     noise = (noise - 0.5) * 0.18
  398.  
  399.     q = baseline * (0.5 + 0.5 * inland_decay) * (0.6 + 0.4 * elev_penalty) + noise
  400.     q = np.clip(q, 0, 1)
  401.     return q
  402.  
  403. def compute_rainfall(u, v, height_norm, humidity, temperature_C, water_mask, seed=42):
  404.     """
  405.    Rainfall rate (mm/hr): orographic on land + convective everywhere.
  406.    """
  407.     H, W = height_norm.shape
  408.     dy = 1.0 / max(H - 1, 1)
  409.     dx = 1.0 / max(W - 1, 1)
  410.     dzdy = np.gradient(height_norm, dy, axis=0)
  411.     dzdx = np.gradient(height_norm, dx, axis=1)
  412.  
  413.     upslope = u * dzdx + v * dzdy
  414.     orog = np.clip(upslope, 0, None)
  415.     orog = orog * (~water_mask)  # suppress orographic rain over oceans
  416.  
  417.     t_excess = np.clip(temperature_C - (np.quantile(temperature_C, 0.6)), 0, None)
  418.     conv_noise = fbm_noise((H, W), octaves=3, base_sigma=min(H, W)/18, seed=seed+404)
  419.     conv = t_excess * humidity * conv_noise
  420.  
  421.     rain = 30.0 * (0.65 * orog + 0.35 * conv)
  422.     rain = gaussian_filter(rain, sigma=0.8)
  423.     return rain
  424.  
  425. def wind_speed(u, v):
  426.     return np.sqrt(u*u + v*v)
  427.  
  428. def vorticity(u, v):
  429.     H, W = u.shape
  430.     dy = 1.0 / max(H - 1, 1)
  431.     dx = 1.0 / max(W - 1, 1)
  432.     dvdx = np.gradient(v, dx, axis=1)
  433.     dudy = np.gradient(u, dy, axis=0)
  434.     return dvdx - dudy
  435.  
  436. def find_local_extrema(field, num_max=3, num_min=3, neighborhood=21):
  437.     f = gaussian_filter(field, sigma=2.0)
  438.     max_mask = f == maximum_filter(f, size=neighborhood, mode='reflect')
  439.     max_coords = np.argwhere(max_mask)
  440.     max_vals = f[max_mask]
  441.     if len(max_vals) > 0:
  442.         idx = np.argsort(max_vals)[::-1][:num_max]
  443.         hot_points = max_coords[idx]
  444.     else:
  445.         hot_points = np.empty((0, 2), dtype=int)
  446.     g = -f
  447.     min_mask = g == maximum_filter(g, size=neighborhood, mode='reflect')
  448.     min_coords = np.argwhere(min_mask)
  449.     min_vals = f[min_mask]
  450.     if len(min_vals) > 0:
  451.         jdx = np.argsort(min_vals)[:num_min]
  452.         cold_points = min_coords[jdx]
  453.     else:
  454.         cold_points = np.empty((0, 2), dtype=int)
  455.     return hot_points, cold_points
  456.  
  457. def find_storm_centers(u, v, num_centers=3, neighborhood=31, min_quantile=0.92):
  458.     vort = vorticity(u, v)
  459.     vort_s = gaussian_filter(vort, sigma=2.0)
  460.     thresh = np.quantile(vort_s, min_quantile)
  461.     mask = (vort_s >= thresh) & (vort_s == maximum_filter(vort_s, size=neighborhood))
  462.     coords = np.argwhere(mask)
  463.     vals = vort_s[mask]
  464.     if len(vals) > 0:
  465.         idx = np.argsort(vals)[::-1][:num_centers]
  466.         pts = coords[idx]
  467.     else:
  468.         pts = np.empty((0, 2), dtype=int)
  469.     return pts, vort_s
  470.  
  471.  
  472. # -----------------------------
  473. # Colormaps and drawing helpers
  474. # -----------------------------
  475.  
  476. def _get_cmap(name: str):
  477.     cmaps = getattr(mpl, "colormaps", None)
  478.     if cmaps is not None:
  479.         return cmaps.get_cmap(name)
  480.     # Fallback for older Matplotlib
  481.     from matplotlib import cm as _cm  # local import to avoid Pylance noise
  482.     return _cm.get_cmap(name)
  483.  
  484. def build_temperature_cmap():
  485.     return _get_cmap('coolwarm')
  486.  
  487. def build_rain_cmap():
  488.     base = _get_cmap('Blues')
  489.     colors = np.array(base(np.linspace(0, 1, 256)), copy=True)  # ensure writable ndarray
  490.     colors[:, -1] = np.linspace(0.0, 0.85, 256)
  491.     return ListedColormap(colors)
  492.  
  493. def annotate_points(ax, pts, labels, color='k'):
  494.     for (y, x), label in zip(pts, labels):
  495.         ax.text(
  496.             x, y, label,
  497.             color='white',
  498.             ha='center', va='center',
  499.             fontsize=10, weight='bold',
  500.             bbox=dict(boxstyle='round,pad=0.25', fc=color, ec='none', alpha=0.8),
  501.             zorder=10,
  502.             path_effects=[pe.withStroke(linewidth=1.5, foreground='black', alpha=0.3)]
  503.         )
  504.  
  505. def draw_fronts(ax, T):
  506.     gy, gx = np.gradient(gaussian_filter(T, sigma=1.2))
  507.     gradmag = np.hypot(gx, gy)
  508.     level = np.quantile(gradmag, 0.93)
  509.     cs = ax.contour(gradmag, levels=[level], colors=['magenta'], linewidths=1.0, linestyles='dashed', alpha=0.7)
  510.     return cs
  511.  
  512.  
  513. # -----------------------------
  514. # Phase 1: Static visualization
  515. # -----------------------------
  516.  
  517. def render_static_map(texture, T, rain, u, v, out_path, dpi=200, show=False,
  518.                       draw_quiver=False, title='Phase 1: Static Weather', pb=None):
  519.     if pb:
  520.         pb.new_task("Rendering static map", total=3)
  521.  
  522.     H, W, _ = texture.shape
  523.     aspect = W / H
  524.     fig_w = 12
  525.     fig_h = max(6, fig_w / aspect)
  526.  
  527.     fig, ax = plt.subplots(figsize=(fig_w, fig_h), dpi=dpi)
  528.     ax.set_axis_off()
  529.  
  530.     ax.imshow(texture, interpolation='bilinear', zorder=0)
  531.     imT = ax.imshow(T, cmap=build_temperature_cmap(), alpha=0.35, zorder=2)
  532.  
  533.     t_min, t_max = float(np.nanmin(T)), float(np.nanmax(T))
  534.     levels = np.linspace(t_min, t_max, 12)
  535.     cs = ax.contour(T, levels=levels, colors=['k'], linewidths=0.7, alpha=0.7, zorder=3)
  536.     ax.clabel(cs, inline=1, fontsize=8, fmt='%1.0f°C')
  537.  
  538.     imR = ax.imshow(rain, cmap=build_rain_cmap(), alpha=0.75, zorder=4)
  539.  
  540.     ds = max(1, int(min(H, W) / 120))
  541.     YY, XX = np.mgrid[0:H:ds, 0:W:ds]
  542.     U = u[::ds, ::ds]
  543.     V = v[::ds, ::ds]
  544.     sp = wind_speed(U, V)
  545.     lw = 1.5 * (sp / (sp.max() + 1e-9)) + 0.5
  546.     ax.streamplot(XX, YY, U, V, color='white', linewidth=lw, density=1.3, arrowsize=0.8, minlength=0.15, zorder=5)
  547.  
  548.     if draw_quiver:
  549.         step = ds * 3
  550.         ax.quiver(np.arange(0, W, step), np.arange(0, H, step),
  551.                   u[::step, ::step], v[::step, ::step],
  552.                   color='white', alpha=0.8, zorder=6, pivot='mid', scale=50)
  553.  
  554.     hot_pts, cold_pts = find_local_extrema(T, num_max=3, num_min=3, neighborhood=31)
  555.     if len(hot_pts) > 0:
  556.         annotate_points(ax, hot_pts, ['HOT'] * len(hot_pts), color='#d62728')
  557.     if len(cold_pts) > 0:
  558.         annotate_points(ax, cold_pts, ['COLD'] * len(cold_pts), color='#1f77b4')
  559.  
  560.     storm_pts, _ = find_storm_centers(u, v, num_centers=3, neighborhood=31, min_quantile=0.92)
  561.     if len(storm_pts) > 0:
  562.         annotate_points(ax, storm_pts, ['L'] * len(storm_pts), color='#2ca02c')
  563.  
  564.     draw_fronts(ax, T)
  565.  
  566.     if pb:
  567.         pb.advance(1)
  568.  
  569.     cbar_T = fig.colorbar(imT, ax=ax, fraction=0.046, pad=0.02)
  570.     cbar_T.set_label('Temperature (°C)')
  571.     cbar_R = fig.colorbar(imR, ax=ax, fraction=0.046, pad=0.02)
  572.     cbar_R.set_label('Rainfall (mm/hr)')
  573.  
  574.     ax.set_title(title, fontsize=14, weight='bold')
  575.     fig.tight_layout()
  576.  
  577.     if pb:
  578.         pb.advance(1)
  579.  
  580.     fig.savefig(out_path, bbox_inches='tight', dpi=dpi)
  581.     if show:
  582.         plt.show()
  583.     plt.close(fig)
  584.  
  585.     if pb:
  586.         pb.done("Static map saved")
  587.  
  588.  
  589. # -----------------------------
  590. # Phase 2: Animation (wind/rain)
  591. # -----------------------------
  592.  
  593. def bilinear_sample(a, x, y):
  594.     H, W = a.shape
  595.     x = np.clip(x, 0.0, W - 1.000001)
  596.     y = np.clip(y, 0.0, H - 1.000001)
  597.     x0 = np.floor(x).astype(np.int32)
  598.     y0 = np.floor(y).astype(np.int32)
  599.     x1 = np.clip(x0 + 1, 0, W - 1)
  600.     y1 = np.clip(y0 + 1, 0, H - 1)
  601.     fx = x - x0
  602.     fy = y - y0
  603.     v00 = a[y0, x0]
  604.     v10 = a[y0, x1]
  605.     v01 = a[y1, x0]
  606.     v11 = a[y1, x1]
  607.     v0 = v00 * (1.0 - fx) + v10 * fx
  608.     v1 = v01 * (1.0 - fx) + v11 * fx
  609.     return v0 * (1.0 - fy) + v1 * fy
  610.  
  611. def advect(x, y, u, v, px_per_frame):
  612.     ux = bilinear_sample(u, x, y) * px_per_frame
  613.     vy = bilinear_sample(v, x, y) * px_per_frame
  614.     return x + ux, y + vy, np.sqrt(ux*ux + vy*vy)
  615.  
  616. def wrap_or_respawn(x, y, W, H, rng, respawn_mask=None, edge='wrap'):
  617.     if edge == 'wrap':
  618.         x = np.mod(x, W)
  619.         y = np.mod(y, H)
  620.         return x, y
  621.     need = respawn_mask if respawn_mask is not None else (
  622.         (x < 0) | (x >= W) | (y < 0) | (y >= H)
  623.     )
  624.     n = np.count_nonzero(need)
  625.     if n > 0:
  626.         x[need] = rng.uniform(0, W, size=n)
  627.         y[need] = rng.uniform(0, H, size=n)
  628.     return x, y
  629.  
  630. def sample_positions_from_weight_map(weight_map, n, rng):
  631.     H, W = weight_map.shape
  632.     w = weight_map.ravel().astype(np.float64)
  633.     s = w.sum()
  634.     if s <= 0:
  635.         x = rng.uniform(0, W, size=n)
  636.         y = rng.uniform(0, H, size=n)
  637.         return x, y
  638.     p = w / s
  639.     idx = rng.choice(W * H, size=n, replace=True, p=p)
  640.     y_idx, x_idx = np.divmod(idx, W)
  641.     x = x_idx + rng.uniform(0, 1, size=n)
  642.     y = y_idx + rng.uniform(0, 1, size=n)
  643.     return x, y
  644.  
  645. def animate_weather(texture, T, rain_field, u, v,
  646.                     out_path='weather_anim.mp4', mode='both',
  647.                     n_wind=2000, n_rain=2000, fps=30, duration=32,
  648.                     dpi=160, show=False, seed=42, draw_isotherms=True,
  649.                     temp_alpha=0.28, rain_alpha=0.35, streamlines=False,
  650.                     blit=True, flow_scale=0.7, base_title='Phase 2: Animated Weather',
  651.                     pb=None):
  652.     rng = np.random.default_rng(seed)
  653.     H, W, _ = texture.shape
  654.     frames = int(max(1, duration) * max(1, fps))
  655.  
  656.     spd = wind_speed(u, v)
  657.     ref = np.percentile(spd, 99.0) + 1e-9
  658.     px_per_frame = (flow_scale * min(H, W)) / (ref * fps)
  659.  
  660.     aspect = W / H
  661.     fig_w = 12
  662.     fig_h = max(6, fig_w / aspect)
  663.     fig, ax = plt.subplots(figsize=(fig_w, fig_h), dpi=dpi)
  664.     ax.set_axis_off()
  665.     ax.set_xlim(0, W)
  666.     ax.set_ylim(H, 0)
  667.  
  668.     ax.imshow(texture, interpolation='bilinear', zorder=0)
  669.     ax.imshow(T, cmap=build_temperature_cmap(), alpha=temp_alpha, zorder=1)
  670.  
  671.     if draw_isotherms:
  672.         levels = np.linspace(float(np.nanmin(T)), float(np.nanmax(T)), 10)
  673.         cs = ax.contour(T, levels=levels, colors=['k'], linewidths=0.6, alpha=0.6, zorder=2)
  674.         ax.clabel(cs, inline=1, fontsize=8, fmt='%1.0f°C')
  675.  
  676.     r_cmap = build_rain_cmap()
  677.     ax.imshow(rain_field, cmap=r_cmap, alpha=rain_alpha, zorder=3)
  678.  
  679.     if streamlines:
  680.         ds = max(1, int(min(H, W) / 120))
  681.         YY, XX = np.mgrid[0:H:ds, 0:W:ds]
  682.         U = u[::ds, ::ds]
  683.         V = v[::ds, ::ds]
  684.         sp = wind_speed(U, V)
  685.         lw = 1.2 * (sp / (sp.max() + 1e-9)) + 0.4
  686.         ax.streamplot(XX, YY, U, V, color='white', linewidth=lw, density=1.2, arrowsize=0.7, minlength=0.15, zorder=4)
  687.  
  688.     artists = []
  689.  
  690.     # Wind streaks
  691.     lc = None
  692.     wind_x = wind_y = None
  693.     wind_norm = Normalize(vmin=0.0, vmax=np.percentile(spd, 95.0))
  694.     if mode in ('wind', 'both') and n_wind > 0:
  695.         wind_x = rng.uniform(0, W, size=n_wind).astype(np.float32)
  696.         wind_y = rng.uniform(0, H, size=n_wind).astype(np.float32)
  697.         lc = LineCollection(segments=[], cmap=_get_cmap('plasma'), norm=wind_norm,
  698.                             linewidths=1.2, alpha=0.9, zorder=6)
  699.         lc.set_array(np.zeros(n_wind, dtype=np.float32))
  700.         if blit:
  701.             lc.set_animated(True)
  702.         ax.add_collection(lc)  # type: ignore[arg-type]
  703.         artists.append(lc)
  704.  
  705.     # Rain scatter (init arrays so Pylance knows they're arrays)
  706.     rain_x = np.empty(0, dtype=np.float32)
  707.     rain_y = np.empty(0, dtype=np.float32)
  708.     rain_alpha_arr = np.empty(0, dtype=np.float32)
  709.     rain_life = np.empty(0, dtype=np.float32)
  710.     sc = None
  711.     rain_map = rain_field - rain_field.min()
  712.     if rain_map.max() > 0:
  713.         rain_map = rain_map / (rain_map.max() + 1e-9)
  714.     max_rain = max(0, n_rain)
  715.     if mode in ('rain', 'both') and max_rain > 0:
  716.         init_n = min(max_rain, max(100, max_rain // 5))
  717.         rx, ry = sample_positions_from_weight_map(rain_map, init_n, rng)
  718.         rain_x = rx.astype(np.float32)
  719.         rain_y = ry.astype(np.float32)
  720.         rain_alpha_arr = np.ones_like(rain_x, dtype=np.float32) * 0.85
  721.         rain_life = rng.uniform(0.8 * fps, 2.0 * fps, size=rain_x.shape[0]).astype(np.float32)
  722.         sc = ax.scatter(rain_x, rain_y, s=6.0, c='white', alpha=rain_alpha_arr, edgecolors='none', zorder=7)
  723.         if blit:
  724.             sc.set_animated(True)
  725.         artists.append(sc)
  726.  
  727.     ax.set_title(f"{base_title}", fontsize=14, weight='bold')
  728.     fig.canvas.draw()
  729.  
  730.     def init():
  731.         return artists
  732.  
  733.     def update(frame_idx):
  734.         nonlocal rain_x, rain_y, rain_alpha_arr, rain_life
  735.  
  736.         if lc is not None:
  737.             new_x, new_y, _ = advect(wind_x, wind_y, u, v, px_per_frame)
  738.             segs = np.stack([np.stack([wind_x, wind_y], axis=1),
  739.                              np.stack([new_x, new_y], axis=1)], axis=1)
  740.             lc.set_segments(segs)
  741.             s_here = bilinear_sample(spd, wind_x, wind_y)
  742.             lc.set_array(s_here.astype(np.float32))
  743.             wind_x[:] = new_x
  744.             wind_y[:] = new_y
  745.             wind_x[:], wind_y[:] = wrap_or_respawn(wind_x, wind_y, W, H, rng, edge='wrap')
  746.  
  747.         if sc is not None:
  748.             spawn_per_frame = max(1, int(0.03 * max_rain))
  749.             rx, ry = sample_positions_from_weight_map(rain_map, spawn_per_frame, rng)
  750.             if rain_x.size == 0:
  751.                 rain_x = rx.astype(np.float32)
  752.                 rain_y = ry.astype(np.float32)
  753.                 rain_alpha_arr = np.ones_like(rain_x) * 0.85
  754.                 rain_life = rng.uniform(0.8 * fps, 2.0 * fps, size=rain_x.shape[0]).astype(np.float32)
  755.             else:
  756.                 rain_x = np.concatenate([rain_x, rx.astype(np.float32)])
  757.                 rain_y = np.concatenate([rain_y, ry.astype(np.float32)])
  758.                 rain_alpha_arr = np.concatenate([rain_alpha_arr, np.ones(rx.shape[0], dtype=np.float32) * 0.85])
  759.                 rain_life = np.concatenate([rain_life, rng.uniform(0.8 * fps, 2.0 * fps, size=rx.shape[0]).astype(np.float32)])
  760.  
  761.             if rain_x.size > max_rain:
  762.                 rain_x = rain_x[-max_rain:]
  763.                 rain_y = rain_y[-max_rain:]
  764.                 rain_alpha_arr = rain_alpha_arr[-max_rain:]
  765.                 rain_life = rain_life[-max_rain:]
  766.  
  767.             rain_x, rain_y, _ = advect(rain_x, rain_y, u, v, px_per_frame)
  768.             rain_life -= 1.0
  769.             rain_alpha_arr *= 0.985
  770.             out = (rain_x < -2) | (rain_x > W + 2) | (rain_y < -2) | (rain_y > H + 2) | (rain_life <= 0) | (rain_alpha_arr < 0.05)
  771.             if np.any(out):
  772.                 keep_mask = ~out
  773.                 rain_x = rain_x[keep_mask]
  774.                 rain_y = rain_y[keep_mask]
  775.                 rain_alpha_arr = rain_alpha_arr[keep_mask]
  776.                 rain_life = rain_life[keep_mask]
  777.  
  778.             sc.set_offsets(np.c_[rain_x, rain_y])
  779.             sc.set_alpha(rain_alpha_arr)
  780.  
  781.         return artists
  782.  
  783.     anim = animation.FuncAnimation(fig, update, init_func=init, frames=frames, interval=1000.0 / fps, blit=blit)
  784.  
  785.     # Progress bar for saving frames
  786.     if out_path:
  787.         ext = os.path.splitext(out_path)[1].lower()
  788.         total_frames = frames
  789.         if pb:
  790.             pb.new_task("Exporting animation frames", total=total_frames)
  791.  
  792.         def progress_callback(i, n):
  793.             if pb:
  794.                 pb.update(i + 1)
  795.  
  796.         try:
  797.             if ext == '.mp4':
  798.                 Writer = animation.FFMpegWriter
  799.                 writer = Writer(fps=fps, metadata=dict(artist='You'), bitrate=6000)
  800.                 anim.save(out_path, writer=writer, dpi=dpi, progress_callback=progress_callback)
  801.             elif ext == '.gif':
  802.                 anim.save(out_path, writer='pillow', dpi=dpi, progress_callback=progress_callback)
  803.             else:
  804.                 Writer = animation.FFMpegWriter
  805.                 writer = Writer(fps=fps, metadata=dict(artist='You'), bitrate=6000)
  806.                 anim.save(out_path if ext else out_path + '.mp4', writer=writer, dpi=dpi, progress_callback=progress_callback)
  807.         except Exception as e:
  808.             if pb:
  809.                 pb.done("Export failed")
  810.             print("\nFailed to save animation:", e)
  811.             print("Tip: for MP4, ensure ffmpeg is installed and on PATH.")
  812.         else:
  813.             if pb:
  814.                 pb.done("Animation saved")
  815.  
  816.     if show:
  817.         plt.show()
  818.     plt.close(fig)
  819.  
  820.  
  821. # -----------------------------
  822. # Phase 3: 3D scene
  823. # -----------------------------
  824.  
  825. def render_3d_matplotlib(texture, height_norm, T, rain, u, v, out_path=None, show=False,
  826.                          max_elev_m=3000.0, title='Phase 3: 3D Weather (Matplotlib)', pb=None):
  827.     if pb:
  828.         pb.new_task("Preparing 3D data", total=3)
  829.  
  830.     H, W = height_norm.shape
  831.     target = 220
  832.     stride = max(1, int(min(H, W) / target))
  833.     h_ds = height_norm[::stride, ::stride]
  834.     T_ds = T[::stride, ::stride]
  835.     rain_ds = rain[::stride, ::stride]
  836.     u_ds = u[::stride, ::stride]
  837.     v_ds = v[::stride, ::stride]
  838.     Hd, Wd = h_ds.shape
  839.  
  840.     if pb:
  841.         pb.advance(1)
  842.  
  843.     Xg, Yg = np.meshgrid(np.linspace(0, 1, Wd), np.linspace(0, 1, Hd))
  844.     Z = h_ds * max_elev_m
  845.  
  846.     t_cmap = build_temperature_cmap()
  847.     t_norm = (T_ds - T_ds.min()) / (T_ds.max() - T_ds.min() + 1e-9)
  848.     face_colors = t_cmap(t_norm)
  849.     rain_norm = np.clip(rain_ds / (np.percentile(rain_ds, 98) + 1e-9), 0, 1)
  850.     blue_tint = np.zeros_like(face_colors)
  851.     blue_tint[..., 2] = 1.0
  852.     alpha_rain = 0.35 * rain_norm
  853.     face_colors = (1 - alpha_rain[..., None]) * face_colors + alpha_rain[..., None] * blue_tint
  854.  
  855.     if pb:
  856.         pb.advance(1)
  857.  
  858.     fig = plt.figure(figsize=(12, 8))
  859.     ax = fig.add_subplot(111, projection='3d')
  860.     ax.plot_surface(Xg, Yg, Z, facecolors=face_colors, rstride=1, cstride=1, linewidth=0.0, antialiased=False, shade=False)
  861.     ax.set_title(title, fontsize=14, weight='bold')
  862.     ax.set_xlabel('X')
  863.     ax.set_ylabel('Y')
  864.     ax.set_zlabel('Elevation (m)')
  865.     ax.view_init(elev=55, azim=-60)
  866.  
  867.     step = max(1, int(min(Hd, Wd) / 22))
  868.     Xq = Xg[::step, ::step]
  869.     Yq = Yg[::step, ::step]
  870.     Zq = Z[::step, ::step] + 30.0
  871.     Uq = u_ds[::step, ::step]
  872.     Vq = v_ds[::step, ::step]
  873.     Wq = np.zeros_like(Uq)
  874.     mag = np.sqrt(Uq**2 + Vq**2) + 1e-9
  875.     Uq = Uq / mag
  876.     Vq = Vq / mag
  877.     ax.quiver(Xq, Yq, Zq, Uq, Vq, Wq, length=0.03, color='white', linewidth=0.6, normalize=False)
  878.  
  879.     fig.tight_layout()
  880.  
  881.     if pb:
  882.         pb.advance(1)
  883.  
  884.     if out_path:
  885.         fig.savefig(out_path, dpi=180, bbox_inches='tight')
  886.  
  887.     if show:
  888.         plt.show()
  889.     plt.close(fig)
  890.  
  891.     if pb:
  892.         pb.done("3D scene saved" if out_path else "3D scene ready")
  893.  
  894. def render_3d_plotly(texture, height_norm, T, rain, out_path=None,
  895.                      max_elev_m=3000.0, title='Phase 3: 3D Weather (Plotly)', pb=None):
  896.     try:
  897.         import plotly.graph_objects as go  # type: ignore
  898.     except Exception:
  899.         print("Plotly not installed. Falling back to Matplotlib 3D.")
  900.         return False
  901.  
  902.     if pb:
  903.         pb.new_task("Preparing 3D (Plotly)", total=2)
  904.  
  905.     H, W = height_norm.shape
  906.     target = 300
  907.     stride = max(1, int(min(H, W) / target))
  908.     h_ds = height_norm[::stride, ::stride]
  909.     T_ds = T[::stride, ::stride]
  910.     Hd, Wd = h_ds.shape
  911.  
  912.     Z = h_ds * max_elev_m
  913.     t_norm = (T_ds - T_ds.min()) / (T_ds.max() - T_ds.min() + 1e-9)
  914.  
  915.     if pb:
  916.         pb.advance(1)
  917.  
  918.     fig = go.Figure(data=[
  919.         go.Surface(
  920.             z=Z,
  921.             x=np.linspace(0, 1, Wd),
  922.             y=np.linspace(0, 1, Hd),
  923.             surfacecolor=t_norm,
  924.             colorscale='RdBu',
  925.             reversescale=True,
  926.             cmin=0.0, cmax=1.0,
  927.             showscale=True,
  928.             contours={"z": {"show": False}},
  929.         )
  930.     ])
  931.     fig.update_layout(
  932.         title=title,
  933.         scene=dict(
  934.             xaxis_title='X', yaxis_title='Y', zaxis_title='Elevation (m)',
  935.             aspectmode='data',
  936.         ),
  937.         width=1000, height=700
  938.     )
  939.  
  940.     if out_path:
  941.         if out_path.lower().endswith('.html'):
  942.             fig.write_html(out_path)
  943.         else:
  944.             try:
  945.                 fig.write_image(out_path)  # requires kaleido
  946.             except Exception:
  947.                 print("Plotly static image export failed (install 'kaleido'). Writing HTML instead.")
  948.                 fig.write_html(os.path.splitext(out_path)[0] + ".html")
  949.     else:
  950.         fig.show()
  951.  
  952.     if pb:
  953.         pb.done("3D scene saved")
  954.  
  955.     return True
  956.  
  957.  
  958. # -----------------------------
  959. # Pipeline orchestration
  960. # -----------------------------
  961.  
  962. def build_fields(heightmap_path, texture_path, sea_level, seed, base_wind_dir, base_wind_speed, max_elev_m, pb=None):
  963.     if pb:
  964.         pb.new_task("Loading inputs", total=8)
  965.  
  966.     h_raw = read_image_as_array(heightmap_path)
  967.     if pb: pb.advance(1)
  968.  
  969.     h_gray = ensure_gray2d(h_raw)
  970.     if pb: pb.advance(1)
  971.  
  972.     height_norm = to_float01(h_gray)
  973.     if pb: pb.advance(1)
  974.  
  975.     tex_raw = read_image_as_array(texture_path, mode='RGB')
  976.     texture = to_float01(tex_raw)
  977.     if pb: pb.advance(1)
  978.  
  979.     if texture.shape[:2] != height_norm.shape:
  980.         texture = resize_to_match(texture, height_norm.shape)
  981.     if pb: pb.advance(1)
  982.  
  983.     H, W = height_norm.shape
  984.     X, Y = compute_grids(H, W)
  985.     if pb: pb.advance(1)
  986.  
  987.     water_mask, coast_dist = detect_water_mask_combined(texture, height_norm, sea_level=sea_level)
  988.     if pb: pb.advance(1)
  989.  
  990.     T = compute_temperature(height_norm, X, Y, water_mask, coast_dist,
  991.                             seed=seed, max_elev_m=max_elev_m)
  992.     if pb: pb.advance(1)
  993.  
  994.     if pb: pb.new_task("Building wind/humidity/rain", total=3)
  995.     u, v = compute_streamfunction(height_norm, X, Y, water_mask,
  996.                                   base_wind_dir_deg=base_wind_dir,
  997.                                   base_wind_speed=base_wind_speed,
  998.                                   seed=seed, k_orog=0.9, k_noise=1.0)
  999.     if pb: pb.advance(1)
  1000.  
  1001.     q = compute_humidity(height_norm, water_mask, coast_dist, seed=seed)
  1002.     if pb: pb.advance(1)
  1003.  
  1004.     rain = compute_rainfall(u, v, height_norm, q, T, water_mask, seed=seed)
  1005.     if pb: pb.advance(1)
  1006.     if pb: pb.done("Fields ready")
  1007.  
  1008.     return texture, height_norm, water_mask, T, rain, u, v
  1009.  
  1010.  
  1011. # -----------------------------
  1012. # Main / CLI
  1013. # -----------------------------
  1014.  
  1015. def main():
  1016.     ap = argparse.ArgumentParser(description="All-in-one Visual Weather Engine (Phases 1–3)")
  1017.     ap.add_argument('--phase', type=int, default=None, choices=[1, 2, 3], help='1=Static, 2=Animation, 3=3D (omit to pick interactively)')
  1018.     ap.add_argument('--heightmap', type=str, default=None, help='Path to heightmap (grayscale or RGB).')
  1019.     ap.add_argument('--texture', type=str, default=None, help='Path to texture map (RGB).')
  1020.     ap.add_argument('--sea-level', type=float, default=None, help='Sea level threshold in normalized height units [0..1], e.g., 0.22')
  1021.     ap.add_argument('--out', type=str, default=None, help='Output file (png/mp4/gif/html). For --phase all, used as base name.')
  1022.     ap.add_argument('--seed', type=int, default=42, help='Random seed.')
  1023.     ap.add_argument('--base-wind-dir', type=float, default=35.0, help='Wind direction deg (0=E, 90=N).')
  1024.     ap.add_argument('--base-wind-speed', type=float, default=6.0, help='Base wind speed (arbitrary units).')
  1025.     ap.add_argument('--max-elev', type=float, default=3000.0, help='Max elevation represented by heightmap (m).')
  1026.  
  1027.     # Phase 1 options
  1028.     ap.add_argument('--dpi', type=int, default=200, help='DPI for static or animation frames.')
  1029.     ap.add_argument('--show', action='store_true', help='Show window after rendering.')
  1030.     ap.add_argument('--draw-quiver', action='store_true', help='Add quiver arrows to static map.')
  1031.  
  1032.     # Phase 2 options
  1033.     ap.add_argument('--fps', type=int, default=30, help='Frames per second (animation).')
  1034.     ap.add_argument('--duration', type=float, default=10.0, help='Duration in seconds (animation).')
  1035.     ap.add_argument('--mode', type=str, default='both', choices=['wind', 'rain', 'both'], help='Particle mode (animation).')
  1036.     ap.add_argument('--n-wind', type=int, default=2000, help='Number of wind particles.')
  1037.     ap.add_argument('--n-rain', type=int, default=2000, help='Max number of rain particles.')
  1038.     ap.add_argument('--flow-scale', type=float, default=0.7, help='Relative particle speed scaling (~0.1..1.5).')
  1039.     ap.add_argument('--no-blit', action='store_true', help='Disable blitting (slower but more compatible).')
  1040.     ap.add_argument('--no-isos', action='store_true', help='Disable isotherm labels in animation.')
  1041.     ap.add_argument('--temp-alpha', type=float, default=0.28, help='Temperature overlay alpha (animation).')
  1042.     ap.add_argument('--rain-alpha', type=float, default=0.55, help='Rain overlay alpha (animation).')
  1043.     ap.add_argument('--streamlines', action='store_true', help='Draw static streamlines in animation.')
  1044.  
  1045.     # Phase 3 options
  1046.     ap.add_argument('--plotly', action='store_true', help='Use Plotly for 3D (if installed).')
  1047.  
  1048.     args = ap.parse_args()
  1049.  
  1050.     pb = ProgressBar(width=42)
  1051.  
  1052.     # Inputs and sea level
  1053.     heightmap_path = prompt_for_existing_file("Enter heightmap path", default=args.heightmap)
  1054.     texture_path = prompt_for_existing_file("Enter texture map path", default=args.texture)
  1055.     sea_level = prompt_for_float("Enter sea level threshold (normalized 0..1)", default=(args.sea_level if args.sea_level is not None else 0.22), minv=0.0, maxv=1.0)
  1056.  
  1057.     # Interactive phase pick if not provided via CLI
  1058.     phases_to_run = [args.phase] if args.phase is not None else prompt_for_phase(default_choice="1")
  1059.  
  1060.     # Build meteorological fields once
  1061.     texture, height_norm, water_mask, T, rain, u, v = build_fields(
  1062.         heightmap_path, texture_path, sea_level, seed=args.seed,
  1063.         base_wind_dir=args.base_wind_dir, base_wind_speed=args.base_wind_speed,
  1064.         max_elev_m=args.max_elev, pb=pb
  1065.     )
  1066.  
  1067.     # Compute shared timestamp for outputs
  1068.     stamp = _dt.datetime.now().strftime("%Y%m%d_%H%M%S")
  1069.  
  1070.     # Helper to build output names
  1071.     def outputs_for_phase(phase, base_out, use_plotly, stamp):
  1072.         if base_out:
  1073.             base_dir, fn = os.path.split(base_out)
  1074.             base_name, _ext = os.path.splitext(fn)
  1075.             base_dir = base_dir if base_dir else "."
  1076.         else:
  1077.             base_dir, base_name = ".", "weather"
  1078.  
  1079.         if phase == 1:
  1080.             path = os.path.join(base_dir, f"{base_name}_map.png")
  1081.         elif phase == 2:
  1082.             path = os.path.join(base_dir, f"{base_name}_anim.mp4")
  1083.         elif phase == 3:
  1084.             path = os.path.join(base_dir, f"{base_name}_3d.html" if use_plotly else f"{base_name}_3d.png")
  1085.         else:
  1086.             path = os.path.join(base_dir, f"{base_name}.out")
  1087.         # Add shared timestamp
  1088.         return add_timestamp_to_path(path, stamp=stamp)
  1089.  
  1090.     # Run requested phases
  1091.     if phases_to_run == [1]:
  1092.         out1 = outputs_for_phase(1, args.out, args.plotly, stamp)
  1093.         render_static_map(texture, T, rain, u, v, out_path=out1, dpi=args.dpi, show=args.show,
  1094.                           draw_quiver=args.draw_quiver, title='Phase 1: Static Weather Visualization', pb=pb)
  1095.         print(f"Saved: {out1}")
  1096.  
  1097.     elif phases_to_run == [2]:
  1098.         out2 = outputs_for_phase(2, args.out, args.plotly, stamp)
  1099.         animate_weather(
  1100.             texture=texture, T=T, rain_field=rain, u=u, v=v,
  1101.             out_path=out2, mode=args.mode, n_wind=args.n_wind, n_rain=args.n_rain,
  1102.             fps=args.fps, duration=args.duration, dpi=args.dpi, show=args.show,
  1103.             seed=args.seed, draw_isotherms=not args.no_isos, temp_alpha=args.temp_alpha,
  1104.             rain_alpha=args.rain_alpha, streamlines=args.streamlines, blit=not args.no_blit,
  1105.             flow_scale=args.flow_scale, base_title='Phase 2: Animated Weather',
  1106.             pb=pb
  1107.         )
  1108.         print(f"Done. Output: {out2}")
  1109.  
  1110.     elif phases_to_run == [3]:
  1111.         out3 = outputs_for_phase(3, args.out, args.plotly, stamp)
  1112.         used_plotly = False
  1113.         if args.plotly:
  1114.             used_plotly = render_3d_plotly(texture, height_norm, T, rain, out_path=out3,
  1115.                                            max_elev_m=args.max_elev, title='Phase 3: 3D Weather (Plotly)', pb=pb)
  1116.         if not used_plotly:
  1117.             # If we were planning HTML but fell back, adjust extension to PNG + timestamp again
  1118.             if out3.lower().endswith(".html"):
  1119.                 out3 = outputs_for_phase(3, args.out, False, stamp)
  1120.             render_3d_matplotlib(texture, height_norm, T, rain, u, v, out_path=out3, show=args.show,
  1121.                                  max_elev_m=args.max_elev, title='Phase 3: 3D Weather (Matplotlib)', pb=pb)
  1122.         print(f"Saved: {out3}")
  1123.  
  1124.     else:
  1125.         # Run all phases with one shared stamp
  1126.         out1 = outputs_for_phase(1, args.out, args.plotly, stamp)
  1127.         out2 = outputs_for_phase(2, args.out, args.plotly, stamp)
  1128.         out3 = outputs_for_phase(3, args.out, args.plotly, stamp)
  1129.  
  1130.         # Phase 1
  1131.         render_static_map(texture, T, rain, u, v, out_path=out1, dpi=args.dpi, show=False,
  1132.                           draw_quiver=args.draw_quiver, title='Phase 1: Static Weather Visualization', pb=pb)
  1133.         print(f"Saved: {out1}")
  1134.  
  1135.         # Phase 2
  1136.         animate_weather(
  1137.             texture=texture, T=T, rain_field=rain, u=u, v=v,
  1138.             out_path=out2, mode=args.mode, n_wind=args.n_wind, n_rain=args.n_rain,
  1139.             fps=args.fps, duration=args.duration, dpi=args.dpi, show=False,
  1140.             seed=args.seed, draw_isotherms=not args.no_isos, temp_alpha=args.temp_alpha,
  1141.             rain_alpha=args.rain_alpha, streamlines=args.streamlines, blit=not args.no_blit,
  1142.             flow_scale=args.flow_scale, base_title='Phase 2: Animated Weather', pb=pb
  1143.         )
  1144.         print(f"Saved: {out2}")
  1145.  
  1146.         # Phase 3
  1147.         used_plotly = False
  1148.         if args.plotly:
  1149.             used_plotly = render_3d_plotly(texture, height_norm, T, rain, out_path=out3,
  1150.                                            max_elev_m=args.max_elev, title='Phase 3: 3D Weather (Plotly)', pb=pb)
  1151.         if not used_plotly:
  1152.             if out3.lower().endswith(".html"):
  1153.                 out3 = outputs_for_phase(3, args.out, False, stamp)
  1154.             render_3d_matplotlib(texture, height_norm, T, rain, u, v, out_path=out3, show=args.show,
  1155.                                  max_elev_m=args.max_elev, title='Phase 3: 3D Weather (Matplotlib)', pb=pb)
  1156.         print(f"Saved: {out3}")
  1157.  
  1158. if __name__ == '__main__':
  1159.     main()
Tags: python
Advertisement
Add Comment
Please, Sign In to add comment
Public Pastes
We use cookies for various purposes including analytics. By continuing to use Pastebin, you agree to our use of cookies as described in the Cookies Policy.  OK, I Understand
Not a member of Pastebin yet?
Sign Up, it unlocks many cool features!