1use crate::{Error, Result, WindField};
12use nalgebra::Vector3;
13
14const EARTH_RADIUS_M: f64 = 6_371_000.0;
16
17fn horizontal(speed_mps: f64, direction_from_deg: f64) -> Vector3<f64> {
20 let toward = (direction_from_deg + 180.0).to_radians();
21 Vector3::new(speed_mps * toward.cos(), 0.0, speed_mps * toward.sin())
22}
23
24#[derive(Debug, Clone, Copy, PartialEq)]
26pub struct Monitor {
27 pub x_m: f64,
29 pub z_m: f64,
31 pub speed_mps: f64,
33 pub direction_from_deg: f64,
36}
37
38impl Monitor {
39 pub fn at(x_m: f64, z_m: f64, speed_mps: f64, direction_from_deg: f64) -> Self {
41 Self {
42 x_m,
43 z_m,
44 speed_mps,
45 direction_from_deg,
46 }
47 }
48
49 pub fn from_geo(
54 lat_deg: f64,
55 lon_deg: f64,
56 speed_mps: f64,
57 dir_compass_from_deg: f64,
58 ref_lat_deg: f64,
59 ref_lon_deg: f64,
60 north_azimuth_deg: f64,
61 ) -> Self {
62 let (x_m, z_m) = local_en_m(lat_deg, lon_deg, ref_lat_deg, ref_lon_deg);
63 Self {
64 x_m,
65 z_m,
66 speed_mps,
67 direction_from_deg: (dir_compass_from_deg - north_azimuth_deg).rem_euclid(360.0),
68 }
69 }
70}
71
72pub fn local_en_m(lat_deg: f64, lon_deg: f64, ref_lat_deg: f64, ref_lon_deg: f64) -> (f64, f64) {
76 let dlat = (lat_deg - ref_lat_deg).to_radians();
77 let dlon = (lon_deg - ref_lon_deg).to_radians();
78 let east = EARTH_RADIUS_M * dlon * ref_lat_deg.to_radians().cos();
79 let north = EARTH_RADIUS_M * dlat;
80 (east, north)
81}
82
83pub fn monitors_from_geo_csv(
88 text: &str,
89 ref_lat_deg: f64,
90 ref_lon_deg: f64,
91 north_azimuth_deg: f64,
92) -> Result<Vec<Monitor>> {
93 let mut monitors = Vec::new();
94 for (number, line) in text.lines().enumerate() {
95 let line = line.trim();
96 if line.is_empty() || line.starts_with('#') {
97 continue;
98 }
99 let fields: Vec<&str> = line.split(',').collect();
100 if fields.len() < 4 {
101 return Err(Error::invalid(
102 "monitor csv",
103 format!(
104 "line {}: expected lat,lon,speed_ms,dir_compass_deg",
105 number + 1
106 ),
107 ));
108 }
109 let num = |token: &str, what: &str| -> Result<f64> {
110 token.trim().parse().map_err(|_| {
111 Error::invalid(
112 "monitor csv",
113 format!("line {}: {what} '{token}' is not numeric", number + 1),
114 )
115 })
116 };
117 let (lat, lon) = (num(fields[0], "lat")?, num(fields[1], "lon")?);
118 let speed = num(fields[2], "speed")?;
119 let dir = num(fields[3], "direction")?;
120 if !lat.is_finite()
121 || lat.abs() > 90.0
122 || !lon.is_finite()
123 || !speed.is_finite()
124 || speed < 0.0
125 || !dir.is_finite()
126 {
127 return Err(Error::invalid(
128 "monitor csv",
129 format!(
130 "line {}: readings must be finite with nonnegative speed",
131 number + 1
132 ),
133 ));
134 }
135 monitors.push(Monitor::from_geo(
136 lat,
137 lon,
138 speed,
139 dir,
140 ref_lat_deg,
141 ref_lon_deg,
142 north_azimuth_deg,
143 ));
144 }
145 Ok(monitors)
146}
147
148#[derive(Debug, Clone)]
156pub struct HeightGrid {
157 pub origin_x: f64,
159 pub origin_y: f64,
161 pub dx: f64,
163 pub dy: f64,
165 pub ncols: usize,
167 pub nrows: usize,
169 pub heights: Vec<f64>,
171}
172
173impl HeightGrid {
174 pub fn new(
176 origin_x: f64,
177 origin_y: f64,
178 dx: f64,
179 dy: f64,
180 ncols: usize,
181 nrows: usize,
182 heights: Vec<f64>,
183 ) -> Result<Self> {
184 if dx <= 0.0 || dy <= 0.0 || !dx.is_finite() || !dy.is_finite() {
185 return Err(Error::invalid("height grid", "spacing must be positive"));
186 }
187 if ncols < 2 || nrows < 2 {
188 return Err(Error::invalid("height grid", "needs at least 2x2 cells"));
189 }
190 if heights.len() != ncols * nrows {
191 return Err(Error::invalid(
192 "height grid",
193 "height count does not match the grid shape",
194 ));
195 }
196 if !heights.iter().all(|h| h.is_finite()) {
197 return Err(Error::invalid("height grid", "heights must be finite"));
198 }
199 Ok(Self {
200 origin_x,
201 origin_y,
202 dx,
203 dy,
204 ncols,
205 nrows,
206 heights,
207 })
208 }
209
210 pub fn height_at(&self, x: f64, z: f64) -> f64 {
213 let fx = ((x - self.origin_x) / self.dx).clamp(0.0, self.ncols as f64 - 1.0);
214 let fy = ((z - self.origin_y) / self.dy).clamp(0.0, self.nrows as f64 - 1.0);
215 let (j, i) = (fx.floor() as usize, fy.floor() as usize);
216 let (rx, ry) = (fx - j as f64, fy - i as f64);
217 let j1 = (j + 1).min(self.ncols - 1);
218 let i1 = (i + 1).min(self.nrows - 1);
219 let f = |i: usize, j: usize| self.heights[i * self.ncols + j];
220 (1.0 - ry) * ((1.0 - rx) * f(i, j) + rx * f(i, j1))
221 + ry * ((1.0 - rx) * f(i1, j) + rx * f(i1, j1))
222 }
223
224 pub fn from_ascii(text: &str) -> Result<Self> {
230 let mut ncols = None;
231 let mut nrows = None;
232 let mut xll = None;
233 let mut yll = None;
234 let mut cellsize = None;
235 let mut file_rows = Vec::new();
236 for line in text.lines() {
237 let line = line.trim();
238 if line.is_empty() {
239 continue;
240 }
241 let mut parts = line.split_whitespace();
242 match parts.next().unwrap_or("").to_ascii_lowercase().as_str() {
243 "ncols" => ncols = parse_asc_num(parts.next()),
244 "nrows" => nrows = parse_asc_num(parts.next()),
245 "xllcorner" | "xllcenter" => xll = parse_asc_num(parts.next()),
246 "yllcorner" | "yllcenter" => yll = parse_asc_num(parts.next()),
247 "cellsize" => cellsize = parse_asc_num(parts.next()),
248 "nodata_value" => continue,
249 _ => {
250 for token in line.split_whitespace() {
251 if let Ok(v) = token.parse::<f64>() {
252 file_rows.push(v);
253 }
254 }
255 }
256 }
257 }
258 let ncols = ncols.ok_or_else(|| Error::invalid("ascii dem", "missing ncols"))? as usize;
259 let nrows = nrows.ok_or_else(|| Error::invalid("ascii dem", "missing nrows"))? as usize;
260 let cellsize = cellsize.ok_or_else(|| Error::invalid("ascii dem", "missing cellsize"))?;
261 let (xll, yll) = (
262 xll.ok_or_else(|| Error::invalid("ascii dem", "missing xllcorner"))?,
263 yll.ok_or_else(|| Error::invalid("ascii dem", "missing yllcorner"))?,
264 );
265 if file_rows.len() != ncols * nrows {
266 return Err(Error::invalid(
267 "ascii dem",
268 "height count does not match header",
269 ));
270 }
271 let mut heights = Vec::with_capacity(file_rows.len());
274 for row in file_rows.chunks_exact(ncols).rev() {
275 heights.extend_from_slice(row);
276 }
277 Self::new(xll, yll, cellsize, cellsize, ncols, nrows, heights)
278 }
279
280 pub fn from_south_up(
284 origin_x: f64,
285 origin_south_y: f64,
286 dx: f64,
287 dy: f64,
288 ncols: usize,
289 nrows: usize,
290 heights: Vec<f64>,
291 ) -> Result<Self> {
292 Self::new(origin_x, origin_south_y, dx, dy, ncols, nrows, heights)
293 }
294
295 pub fn downsample(&self, target_dx: f64) -> Result<HeightGrid> {
301 if target_dx <= self.dx || target_dx <= 0.0 {
302 return Ok(self.clone());
303 }
304 let step = (target_dx / self.dx).round().max(1.0) as usize;
305 let ncols = (self.ncols / step).max(2);
306 let nrows = (self.nrows / step).max(2);
307 let mut heights = Vec::with_capacity(ncols * nrows);
308 for r in 0..nrows {
309 for c in 0..ncols {
310 let mut sum = 0.0;
311 let mut n = 0usize;
312 for ii in 0..step {
313 for jj in 0..step {
314 let r0 = r * step + ii;
315 let c0 = c * step + jj;
316 if r0 < self.nrows && c0 < self.ncols {
317 sum += self.heights[r0 * self.ncols + c0];
318 n += 1;
319 }
320 }
321 }
322 heights.push(sum / n as f64);
323 }
324 }
325 HeightGrid::new(
326 self.origin_x,
327 self.origin_y,
328 self.dx * step as f64,
329 self.dy * step as f64,
330 ncols,
331 nrows,
332 heights,
333 )
334 }
335}
336
337fn parse_asc_num(token: Option<&str>) -> Option<f64> {
338 token.and_then(|s| s.parse().ok())
339}
340
341type Header = (usize, usize, f64, f64, f64, Option<f64>, bool);
345
346fn parse_header(text: &str) -> Result<Header> {
347 let mut ncols = None;
348 let mut nrows = None;
349 let mut xll = None;
350 let mut yll = None;
351 let mut cellsize = None;
352 let mut nodata = None;
353 let mut msb = true;
354 for line in text.lines() {
355 let mut parts = line.split_whitespace();
356 let key = parts.next().unwrap_or("").to_ascii_lowercase();
357 let val = parts.next();
358 match key.as_str() {
359 "ncols" => ncols = val.and_then(|s| s.parse().ok()),
360 "nrows" => nrows = val.and_then(|s| s.parse().ok()),
361 "xllcorner" | "xllcenter" => xll = val.and_then(|s| s.parse().ok()),
362 "yllcorner" | "yllcenter" => yll = val.and_then(|s| s.parse().ok()),
363 "cellsize" => cellsize = val.and_then(|s| s.parse().ok()),
364 "nodata_value" | "nodata" => nodata = val.and_then(|s| s.parse().ok()),
365 "byteorder" => msb = val.is_none_or(|s| s.eq_ignore_ascii_case("msbfirst")),
366 _ => {}
367 }
368 }
369 let ncols = ncols.ok_or_else(|| Error::invalid("header", "missing ncols"))?;
370 let nrows = nrows.ok_or_else(|| Error::invalid("header", "missing nrows"))?;
371 let cellsize = cellsize.ok_or_else(|| Error::invalid("header", "missing cellsize"))?;
372 let xll = xll.ok_or_else(|| Error::invalid("header", "missing xllcorner"))?;
373 let yll = yll.ok_or_else(|| Error::invalid("header", "missing yllcorner"))?;
374 Ok((ncols, nrows, xll, yll, cellsize, nodata, msb))
375}
376
377pub fn from_flt(flt: &[u8], hdr_text: &str) -> Result<HeightGrid> {
382 let (ncols, nrows, xll, yll, cellsize, nodata, msb) = parse_header(hdr_text)?;
383 let expected = ncols * nrows * 4;
384 if flt.len() != expected {
385 return Err(Error::invalid(
386 "flt",
387 format!("expected {expected} bytes, got {}", flt.len()),
388 ));
389 }
390 let mut heights = Vec::with_capacity(ncols * nrows);
391 for chunk in flt.as_chunks::<4>().0 {
392 let value = if msb {
393 f32::from_be_bytes(*chunk)
394 } else {
395 f32::from_le_bytes(*chunk)
396 };
397 heights.push(f64::from(value));
398 }
399 if let Some(nodata) = nodata {
402 if heights.contains(&nodata) {
403 return Err(Error::invalid(
404 "flt",
405 "raster contains NODATA cells; terrain wind needs a dense DEM",
406 ));
407 }
408 }
409 let mut south_first = Vec::with_capacity(heights.len());
412 for row in heights.chunks_exact(ncols).rev() {
413 south_first.extend_from_slice(row);
414 }
415 HeightGrid::new(xll, yll, cellsize, cellsize, ncols, nrows, south_first)
416}
417
418#[derive(Debug, Clone, Copy, PartialEq, Eq)]
423pub enum SrtmResolution {
424 OneArcSecond,
426 ThreeArcSeconds,
428}
429
430impl SrtmResolution {
431 pub fn samples(self) -> usize {
433 match self {
434 Self::OneArcSecond => 3601,
435 Self::ThreeArcSeconds => 1201,
436 }
437 }
438}
439
440pub fn hgt_grid(
447 data: &[u8],
448 resolution: SrtmResolution,
449 north_lat: f64,
450 west_lon: f64,
451 ref_lat: f64,
452 ref_lon: f64,
453) -> Result<HeightGrid> {
454 let samples = resolution.samples();
455 let expected = samples * samples * 2;
456 if data.len() != expected {
457 return Err(Error::invalid(
458 "hgt",
459 format!("expected {expected} bytes, got {}", data.len()),
460 ));
461 }
462 let mut file_rows = Vec::with_capacity(samples * samples);
464 for chunk in data.as_chunks::<2>().0 {
465 file_rows.push(f64::from(i16::from_be_bytes(*chunk)));
466 }
467 let mut heights = Vec::with_capacity(file_rows.len());
468 for row in file_rows.chunks_exact(samples).rev() {
469 heights.extend_from_slice(row);
470 }
471 let sw = local_en_m(north_lat - 1.0, west_lon, north_lat, west_lon);
474 let ne = local_en_m(north_lat, west_lon + 1.0, north_lat, west_lon);
475 let dx = (ne.0 - sw.0) / (samples - 1) as f64;
476 let dy = (ne.1 - sw.1) / (samples - 1) as f64;
477 let (ref_x, ref_z) = local_en_m(ref_lat, ref_lon, north_lat, west_lon);
478 HeightGrid::new(
479 sw.0 - ref_x,
480 sw.1 - ref_z,
481 dx,
482 dy,
483 samples,
484 samples,
485 heights,
486 )
487}
488
489#[derive(Debug, Clone, Copy, PartialEq)]
491pub struct GeoTransform {
492 pub origin_x: f64,
494 pub origin_y: f64,
497 pub pixel_scale_x: f64,
499 pub pixel_scale_y: f64,
501}
502
503#[cfg(feature = "tiff")]
509pub fn geotransform<R: std::io::Read + std::io::Seek>(
510 decoder: &mut tiff::decoder::Decoder<R>,
511) -> Result<GeoTransform> {
512 if let Ok(transform) = decoder.get_tag_f64_vec(tiff::tags::Tag::ModelTransformationTag) {
514 if transform.len() != 16 {
515 return Err(Error::invalid(
516 "geotiff",
517 "model transform must be 16 values",
518 ));
519 }
520 let is_axis_aligned = [
521 transform[1],
522 transform[2],
523 transform[4],
524 transform[6],
525 transform[8],
526 transform[9],
527 ]
528 .iter()
529 .all(|v| v.abs() < 1e-12);
530 if !is_axis_aligned {
531 return Err(Error::invalid(
532 "geotiff",
533 "rotated model transform is not supported by a HeightGrid",
534 ));
535 }
536 let (sx, sy) = (transform[0], transform[5]);
537 return Ok(GeoTransform {
544 origin_x: transform[3],
545 origin_y: transform[7],
546 pixel_scale_x: sx,
547 pixel_scale_y: -sy,
548 });
549 }
550
551 let tiepoint = decoder
552 .get_tag_f64_vec(tiff::tags::Tag::ModelTiepointTag)
553 .map_err(|_| Error::invalid("geotiff", "missing ModelTiepointTag"))?;
554 let scale = decoder
555 .get_tag_f64_vec(tiff::tags::Tag::ModelPixelScaleTag)
556 .map_err(|_| Error::invalid("geotiff", "missing ModelPixelScaleTag"))?;
557 if scale.len() != 3 {
558 return Err(Error::invalid("geotiff", "pixel scale must be 3 values"));
559 }
560 let (sx, sy) = (scale[0], scale[1]);
561 if sx <= 0.0 {
562 return Err(Error::invalid("geotiff", "pixel scale x must be positive"));
563 }
564 if sy.abs() <= 0.0 {
565 return Err(Error::invalid("geotiff", "pixel scale y must be nonzero"));
566 }
567 let tie_model_x = tiepoint
575 .get(3)
576 .copied()
577 .ok_or_else(|| Error::invalid("geotiff", "tiepoint must have 6 values (I J K X Y Z)"))?;
578 let tie_model_y = *tiepoint
579 .get(4)
580 .ok_or_else(|| Error::invalid("geotiff", "tiepoint must have 6 values (I J K X Y Z)"))?;
581 Ok(GeoTransform {
582 origin_x: tie_model_x,
583 origin_y: tie_model_y,
584 pixel_scale_x: sx,
585 pixel_scale_y: sy.abs(),
586 })
587}
588
589#[cfg(feature = "tiff")]
596pub fn geotiff_grid(data: &[u8]) -> Result<HeightGrid> {
597 let mut decoder = tiff::decoder::Decoder::new(std::io::Cursor::new(data))
598 .map_err(|e| Error::invalid("geotiff", format!("cannot open: {e}")))?;
599 let (width, height) = decoder
600 .dimensions()
601 .map_err(|e| Error::invalid("geotiff", format!("cannot read dimensions: {e}")))?;
602 let (width, height) = (width as usize, height as usize);
603 let values = match decoder
604 .read_image()
605 .map_err(|e| Error::invalid("geotiff", format!("cannot decode: {e}")))?
606 {
607 tiff::decoder::DecodingResult::F32(data) => data.into_iter().map(f64::from).collect(),
608 tiff::decoder::DecodingResult::F64(data) => data,
609 tiff::decoder::DecodingResult::U8(data) => data.into_iter().map(f64::from).collect(),
610 tiff::decoder::DecodingResult::U16(data) => data.into_iter().map(f64::from).collect(),
611 tiff::decoder::DecodingResult::U32(data) => data.into_iter().map(f64::from).collect(),
612 tiff::decoder::DecodingResult::I8(data) => data.into_iter().map(f64::from).collect(),
613 tiff::decoder::DecodingResult::I16(data) => data.into_iter().map(f64::from).collect(),
614 tiff::decoder::DecodingResult::I32(data) => data.into_iter().map(f64::from).collect(),
615 tiff::decoder::DecodingResult::U64(data) => data.into_iter().map(|v| v as f64).collect(),
616 tiff::decoder::DecodingResult::I64(data) => data.into_iter().map(|v| v as f64).collect(),
617 _ => {
618 return Err(Error::invalid(
619 "geotiff",
620 "unsupported sample format (expected integer or float samples)",
621 ))
622 }
623 };
624 if values.len() != width * height {
625 return Err(Error::invalid(
626 "geotiff",
627 "decoded size does not match header",
628 ));
629 }
630 let mut south_first = Vec::with_capacity(values.len());
633 for row in values.chunks_exact(width).rev() {
634 south_first.extend_from_slice(row);
635 }
636 let geot = geotransform(&mut decoder);
639 match geot {
640 Ok(geot) => {
641 let south_y = geot.origin_y - (height as f64 - 1.0) * geot.pixel_scale_y;
644 HeightGrid::new(
645 geot.origin_x,
646 south_y,
647 geot.pixel_scale_x,
648 geot.pixel_scale_y,
649 width,
650 height,
651 south_first,
652 )
653 }
654 Err(_) => HeightGrid::new(0.0, 0.0, 1.0, 1.0, width, height, south_first),
655 }
656}
657
658#[derive(Debug, Clone)]
660pub struct TerrainFlowField {
661 grid: HeightGrid,
662 levels_agl: Vec<f64>,
664 u: Vec<f64>,
667 v: Vec<f64>,
668}
669
670impl TerrainFlowField {
671 pub fn adjust(
679 grid: HeightGrid,
680 base: &dyn WindField,
681 monitors: &[Monitor],
682 levels_agl: &[f64],
683 obs_length_m: f64,
684 surface_decay_m: f64,
685 terrain_scale_m: f64,
686 ) -> Result<Self> {
687 if levels_agl.is_empty() || levels_agl.iter().any(|h| !h.is_finite() || *h < 0.0) {
688 return Err(Error::invalid(
689 "levels_agl",
690 "must be finite and nonnegative",
691 ));
692 }
693 if obs_length_m <= 0.0 || surface_decay_m <= 0.0 || terrain_scale_m <= 0.0 {
694 return Err(Error::invalid("terrain", "length scales must be positive"));
695 }
696 if levels_agl.is_empty() || levels_agl.iter().any(|h| !h.is_finite() || *h < 0.0) {
697 return Err(Error::invalid(
698 "levels_agl",
699 "must be finite and nonnegative",
700 ));
701 }
702 if obs_length_m <= 0.0 || surface_decay_m <= 0.0 {
703 return Err(Error::invalid("terrain", "length scales must be positive"));
704 }
705 for m in monitors {
706 if !m.x_m.is_finite() || !m.z_m.is_finite() || !m.direction_from_deg.is_finite() {
707 return Err(Error::invalid("monitor", "must be finite"));
708 }
709 if !m.speed_mps.is_finite() || m.speed_mps < 0.0 {
710 return Err(Error::invalid("monitor", "speed must be nonnegative"));
711 }
712 }
713 let (ncols, nrows) = (grid.ncols, grid.nrows);
714 let n = ncols * nrows;
715 let inv_dx2 = 1.0 / (grid.dx * grid.dx);
716 let inv_dy2 = 1.0 / (grid.dy * grid.dy);
717 let diag = -2.0 * (inv_dx2 + inv_dy2);
718 let omega = 1.6;
719 let mut u = Vec::with_capacity(levels_agl.len() * n);
720 let mut v = Vec::with_capacity(levels_agl.len() * n);
721 let mut phi = vec![0.0; n];
722 for &level in levels_agl {
723 let mut u0 = vec![0.0; n];
725 let mut v0 = vec![0.0; n];
726 for i in 0..nrows {
727 for j in 0..ncols {
728 let x = grid.origin_x + j as f64 * grid.dx;
729 let z = grid.origin_y + i as f64 * grid.dy;
730 let height = grid.heights[i * ncols + j] + level;
731 let bg = base.wind(0.0, Vector3::new(x, height, z));
732 let (mut cx, mut cz, mut wsum) = (0.0, 0.0, 0.0);
733 for m in monitors {
734 let mw = horizontal(m.speed_mps, m.direction_from_deg);
735 let bg_m = base.wind(
736 0.0,
737 Vector3::new(m.x_m, grid.height_at(m.x_m, m.z_m), m.z_m),
738 );
739 let dx = x - m.x_m;
740 let dz = z - m.z_m;
741 let weight = (-(dx * dx + dz * dz) / (obs_length_m * obs_length_m)).exp()
742 * (-level / surface_decay_m).exp();
743 cx += weight * (mw.x - bg_m.x);
744 cz += weight * (mw.z - bg_m.z);
745 wsum += weight;
746 }
747 if wsum > 0.0 {
748 let w_floor = 0.1;
751 cx /= wsum + w_floor;
752 cz /= wsum + w_floor;
753 }
754 u0[i * ncols + j] = bg.x + cx;
755 v0[i * ncols + j] = bg.z + cz;
756 }
757 }
758 let level_decay = (-level / terrain_scale_m).exp() / terrain_scale_m;
763 let slopes = |u0x: f64, u0z: f64, i: usize, j: usize| -> f64 {
764 let dhdx = (grid.heights[i * ncols + (j + 1).min(ncols - 1)]
765 - grid.heights[i * ncols + j.saturating_sub(1)])
766 / (2.0 * grid.dx);
767 let dhdz = (grid.heights[(i + 1).min(nrows - 1) * ncols + j]
768 - grid.heights[i.saturating_sub(1) * ncols + j])
769 / (2.0 * grid.dy);
770 u0x * dhdx + u0z * dhdz
771 };
772 phi.fill(0.0);
773 for _ in 0..4000 {
774 let mut max_step = 0.0_f64;
775 for i in 1..nrows - 1 {
776 for j in 1..ncols - 1 {
777 let ip = (i + 1).min(nrows - 1);
778 let im = i.saturating_sub(1);
779 let jp = (j + 1).min(ncols - 1);
780 let jm = j.saturating_sub(1);
781 let lap = (phi[i * ncols + jp] + phi[i * ncols + jm]) * inv_dx2
782 + (phi[ip * ncols + j] + phi[im * ncols + j]) * inv_dy2
783 + diag * phi[i * ncols + j];
784 let div = (u0[i * ncols + jp] - u0[i * ncols + jm]) / (2.0 * grid.dx)
785 + (v0[ip * ncols + j] - v0[im * ncols + j]) / (2.0 * grid.dy);
786 let rhs =
787 div - slopes(u0[i * ncols + j], v0[i * ncols + j], i, j) * level_decay;
788 let updated = phi[i * ncols + j] + omega * (rhs - lap) / diag;
789 max_step = max_step.max((updated - phi[i * ncols + j]).abs());
790 phi[i * ncols + j] = updated;
791 }
792 }
793 if max_step < 1e-6 {
794 break;
795 }
796 }
797 for i in 0..nrows {
798 for j in 0..ncols {
799 let edge = i == 0 || i == nrows - 1 || j == 0 || j == ncols - 1;
800 if edge {
801 u.push(u0[i * ncols + j]);
804 v.push(v0[i * ncols + j]);
805 continue;
806 }
807 let jp = (j + 1).min(ncols - 1);
808 let jm = j.saturating_sub(1);
809 let ip = (i + 1).min(nrows - 1);
810 let im = i.saturating_sub(1);
811 let gx = (phi[i * ncols + jp] - phi[i * ncols + jm]) / (2.0 * grid.dx);
812 let gz = (phi[ip * ncols + j] - phi[im * ncols + j]) / (2.0 * grid.dy);
813 u.push(u0[i * ncols + j] - gx);
814 v.push(v0[i * ncols + j] - gz);
815 }
816 }
817 }
818 Ok(Self {
819 grid,
820 levels_agl: levels_agl.to_vec(),
821 u,
822 v,
823 })
824 }
825
826 fn sample_level(&self, level: usize, x: f64, z: f64) -> (f64, f64) {
827 let (ncols, nrows) = (self.grid.ncols, self.grid.nrows);
828 let fx = ((x - self.grid.origin_x) / self.grid.dx).clamp(0.0, ncols as f64 - 1.0);
829 let fy = ((z - self.grid.origin_y) / self.grid.dy).clamp(0.0, nrows as f64 - 1.0);
830 let (j, i) = (fx.floor() as usize, fy.floor() as usize);
831 let (rx, ry) = (fx - j as f64, fy - i as f64);
832 let j1 = (j + 1).min(ncols - 1);
833 let i1 = (i + 1).min(nrows - 1);
834 let base = level * ncols * nrows;
835 let g = |arr: &[f64], i: usize, j: usize| arr[base + i * ncols + j];
836 let sample = |arr: &[f64]| {
837 (1.0 - ry) * ((1.0 - rx) * g(arr, i, j) + rx * g(arr, i, j1))
838 + ry * ((1.0 - rx) * g(arr, i1, j) + rx * g(arr, i1, j1))
839 };
840 (sample(&self.u), sample(&self.v))
841 }
842
843 pub fn horizontal_at(&self, x: f64, y: f64, z: f64) -> (f64, f64) {
845 let agl = y - self.grid.height_at(x, z);
846 let mut level = 0;
847 for (k, &h) in self.levels_agl.iter().enumerate() {
848 if agl <= h {
849 level = k;
850 break;
851 }
852 level = k;
853 }
854 self.sample_level(level, x, z)
855 }
856}
857
858impl WindField for TerrainFlowField {
859 fn wind(&self, _time_s: f64, position_m: Vector3<f64>) -> Vector3<f64> {
860 let (u, v) = self.horizontal_at(position_m.x, position_m.y, position_m.z);
861 Vector3::new(u, 0.0, v)
862 }
863}