veloren_world/sim/
diffusion.rs

1use super::Alt;
2use rayon::prelude::*;
3
4/// From https://github.com/fastscape-lem/fastscapelib-fortran/blob/master/src/Diffusion.f90
5///
6/// See https://fastscape-lem.github.io/fastscapelib-fortran
7///
8/// nx = grid size x
9///
10/// ny = grid size y
11///
12/// xl = x-dimension of the model topography in meters (double precision)
13///
14/// yl = y-dimension of the model topography in meters (double precision)
15///
16/// dt = length of the time step in years (double precision)
17///
18/// ibc = boundary conditions for grid.  For now, we assume all four boundaries
19/// are fixed height, so this parameter is ignored.
20///
21/// h = heights of cells at each cell in the grid.
22///
23/// b = basement height at each cell in the grid (see https://en.wikipedia.org/wiki/Basement_(geology)).
24///
25/// kd = bedrock transport coefficient (or diffusivity) for hillslope processes
26/// in meter squared per year (double precision) at each cell in the grid.
27///
28/// kdsed = sediment transport coefficient (or diffusivity) for hillslope
29/// processes in meter squared per year; (double precision;) note that when
30/// kdsed < 0, its value is not used, i.e., kd for sediment and bedrock have the
31/// same value, regardless of sediment thickness
32/* subroutine Diffusion ()
33
34  ! subroutine to solve the diffusion equation by ADI
35
36  use FastScapeContext
37
38  implicit none
39*/
40#[expect(clippy::needless_late_init)]
41pub fn diffusion(
42    nx: usize,
43    ny: usize,
44    xl: f64,
45    yl: f64,
46    dt: f64,
47    _ibc: (),
48    h: &mut [Alt],
49    b: &mut [Alt],
50    kd: impl Fn(usize) -> f64,
51    kdsed: f64,
52) {
53    let aij = |i: usize, j: usize| j * nx + i;
54    /*
55      double precision, dimension(:), allocatable :: f,diag,sup,inf,res
56      double precision, dimension(:,:), allocatable :: zint,kdint,zintp
57      integer i,j,ij
58      double precision factxp,factxm,factyp,factym,dx,dy
59    */
60    let mut f: Vec<f64>;
61    let mut diag: Vec<f64>;
62    let mut sup: Vec<f64>;
63    let mut inf: Vec<f64>;
64    let mut res: Vec<f64>;
65    let mut zint: Vec<f64>;
66    let mut kdint: Vec<f64>;
67    let mut zintp: Vec<f64>;
68    let mut ij: usize;
69    let mut factxp: f64;
70    let mut factxm: f64;
71    let mut factyp: f64;
72    let mut factym: f64;
73    let dx: f64;
74    let dy: f64;
75    /*
76      character cbc*4
77
78      !print*,'Diffusion'
79
80      write (cbc,'(i4)') ibc
81
82      dx=xl/(nx-1)
83      dy=yl/(ny-1)
84    */
85    // 2048*32/2048/2048
86    // 1 / 64 m
87    dx = xl / /*(nx - 1)*/nx as f64;
88    dy = yl / /*(ny - 1)*/ny as f64;
89    /*
90      ! creates 2D internal arrays to store topo and kd
91
92      allocate (zint(nx,ny),kdint(nx,ny),zintp(nx,ny))
93    */
94    zint = vec![Default::default(); nx * ny];
95    kdint = vec![Default::default(); nx * ny];
96    /*
97      do j=1,ny
98        do i=1,nx
99          ij=(j-1)*nx+i
100          zint(i,j)=h(ij)
101          kdint(i,j)=kd(ij)
102          if (kdsed.gt.0.d0 .and. (h(ij)-b(ij)).gt.1.d-6) kdint(i,j)=kdsed
103        enddo
104      enddo
105
106      zintp = zint
107    */
108    for j in 0..ny {
109        for i in 0..nx {
110            // ij = vec2_as_uniform_idx(i, j);
111            ij = aij(i, j);
112            zint[ij] = h[ij];
113            kdint[ij] = kd(ij);
114            if kdsed > 0.0 && (h[ij] - b[ij]) > 1.0e-6 {
115                kdint[ij] = kdsed;
116            }
117        }
118    }
119
120    zintp = zint.clone();
121    /*
122      ! first pass along the x-axis
123
124      allocate (f(nx),diag(nx),sup(nx),inf(nx),res(nx))
125      f=0.d0
126      diag=0.d0
127      sup=0.d0
128      inf=0.d0
129      res=0.d0
130      do j=2,ny-1
131    */
132    f = vec![0.0; nx];
133    diag = vec![0.0; nx];
134    sup = vec![0.0; nx];
135    inf = vec![0.0; nx];
136    res = vec![0.0; nx];
137    for j in 1..ny - 1 {
138        /*
139            do i=2,nx-1
140            factxp=(kdint(i+1,j)+kdint(i,j))/2.d0*(dt/2.)/dx**2
141            factxm=(kdint(i-1,j)+kdint(i,j))/2.d0*(dt/2.)/dx**2
142              factyp=(kdint(i,j+1)+kdint(i,j))/2.d0*(dt/2.)/dy**2
143            factym=(kdint(i,j-1)+kdint(i,j))/2.d0*(dt/2.)/dy**2
144            diag(i)=1.d0+factxp+factxm
145            sup(i)=-factxp
146            inf(i)=-factxm
147            f(i)=zintp(i,j)+factyp*zintp(i,j+1)-(factyp+factym)*zintp(i,j)+factym*zintp(i,j-1)
148            enddo
149        */
150        for i in 1..nx - 1 {
151            factxp = (kdint[aij(i + 1, j)] + kdint[aij(i, j)]) / 2.0 * (dt / 2.0) / (dx * dx);
152            factxm = (kdint[aij(i - 1, j)] + kdint[aij(i, j)]) / 2.0 * (dt / 2.0) / (dx * dx);
153            factyp = (kdint[aij(i, j + 1)] + kdint[aij(i, j)]) / 2.0 * (dt / 2.0) / (dy * dy);
154            factym = (kdint[aij(i, j - 1)] + kdint[aij(i, j)]) / 2.0 * (dt / 2.0) / (dy * dy);
155            diag[i] = 1.0 + factxp + factxm;
156            sup[i] = -factxp;
157            inf[i] = -factxm;
158            f[i] = zintp[aij(i, j)] + factyp * zintp[aij(i, j + 1)]
159                - (factyp + factym) * zintp[aij(i, j)]
160                + factym * zintp[aij(i, j - 1)];
161        }
162        /*
163        ! left bc
164            if (cbc(4:4).eq.'1') then
165            diag(1)=1.
166            sup(1)=0.
167            f(1)=zintp(1,j)
168            else
169            factxp=(kdint(2,j)+kdint(1,j))/2.d0*(dt/2.)/dx**2
170            factyp=(kdint(1,j+1)+kdint(1,j))/2.d0*(dt/2.)/dy**2
171            factym=(kdint(1,j-1)+kdint(1,j))/2.d0*(dt/2.)/dy**2
172            diag(1)=1.d0+factxp
173            sup(1)=-factxp
174            f(1)=zintp(1,j)+factyp*zintp(1,j+1)-(factyp+factym)*zintp(1,j)+factym*zintp(1,j-1)
175            endif
176        */
177        if true {
178            diag[0] = 1.0;
179            sup[0] = 0.0;
180            f[0] = zintp[aij(0, j)];
181        } else {
182            // reflective boundary
183            factxp = (kdint[aij(1, j)] + kdint[aij(0, j)]) / 2.0 * (dt / 2.0) / (dx * dx);
184            factyp = (kdint[aij(0, j + 1)] + kdint[aij(0, j)]) / 2.0 * (dt / 2.0) / (dy * dy);
185            factym = (kdint[aij(0, j - 1)] + kdint[aij(0, j)]) / 2.0 * (dt / 2.0) / (dy * dy);
186            diag[0] = 1.0 + factxp;
187            sup[0] = -factxp;
188            f[0] = zintp[aij(0, j)] + factyp * zintp[aij(0, j + 1)]
189                - (factyp + factym) * zintp[aij(0, j)]
190                + factym * zintp[aij(0, j - 1)];
191        }
192        /*
193        ! right bc
194            if (cbc(2:2).eq.'1') then
195            diag(nx)=1.
196            inf(nx)=0.
197            f(nx)=zintp(nx,j)
198            else
199            factxm=(kdint(nx-1,j)+kdint(nx,j))/2.d0*(dt/2.)/dx**2
200            factyp=(kdint(nx,j+1)+kdint(nx,j))/2.d0*(dt/2.)/dy**2
201            factym=(kdint(nx,j-1)+kdint(nx,j))/2.d0*(dt/2.)/dy**2
202            diag(nx)=1.d0+factxm
203            inf(nx)=-factxm
204            f(nx)=zintp(nx,j)+factyp*zintp(nx,j+1)-(factyp+factym)*zintp(nx,j)+factym*zintp(nx,j-1)
205            endif
206        */
207        if true {
208            diag[nx - 1] = 1.0;
209            inf[nx - 1] = 0.0;
210            f[nx - 1] = zintp[aij(nx - 1, j)];
211        } else {
212            // reflective boundary
213            factxm = (kdint[aij(nx - 2, j)] + kdint[aij(nx - 1, j)]) / 2.0 * (dt / 2.0) / (dx * dx);
214            factyp =
215                (kdint[aij(nx - 1, j + 1)] + kdint[aij(nx - 1, j)]) / 2.0 * (dt / 2.0) / (dy * dy);
216            factym =
217                (kdint[aij(nx - 1, j - 1)] + kdint[aij(nx - 1, j)]) / 2.0 * (dt / 2.0) / (dy * dy);
218            diag[nx - 1] = 1.0 + factxm;
219            inf[nx - 1] = -factxm;
220            f[nx - 1] = zintp[aij(nx - 1, j)] + factyp * zintp[aij(nx - 1, j + 1)]
221                - (factyp + factym) * zintp[aij(nx - 1, j)]
222                + factym * zintp[aij(nx - 1, j - 1)];
223        }
224        /*
225          call tridag (inf,diag,sup,f,res,nx)
226            do i=1,nx
227            zint(i,j)=res(i)
228            enddo
229        */
230        tridag(&inf, &diag, &sup, &f, &mut res, nx);
231        for i in 0..nx {
232            zint[aij(i, j)] = res[i];
233        }
234        /*
235          enddo
236          deallocate (f,diag,sup,inf,res)
237        */
238    }
239
240    /*
241      ! second pass along y-axis
242
243      allocate (f(ny),diag(ny),sup(ny),inf(ny),res(ny))
244      f=0.d0
245      diag=0.d0
246      sup=0.d0
247      inf=0.d0
248      res=0.d0
249      do i=2,nx-1
250    */
251    f = vec![0.0; ny];
252    diag = vec![0.0; ny];
253    sup = vec![0.0; ny];
254    inf = vec![0.0; ny];
255    res = vec![0.0; ny];
256    for i in 1..nx - 1 {
257        /*
258            do j=2,ny-1
259            factxp=(kdint(i+1,j)+kdint(i,j))/2.d0*(dt/2.)/dx**2
260            factxm=(kdint(i-1,j)+kdint(i,j))/2.d0*(dt/2.)/dx**2
261            factyp=(kdint(i,j+1)+kdint(i,j))/2.d0*(dt/2.)/dy**2
262            factym=(kdint(i,j-1)+kdint(i,j))/2.d0*(dt/2.)/dy**2
263            diag(j)=1.d0+factyp+factym
264            sup(j)=-factyp
265            inf(j)=-factym
266            f(j)=zint(i,j)+factxp*zint(i+1,j)-(factxp+factxm)*zint(i,j)+factxm*zint(i-1,j)
267            enddo
268        */
269        for j in 1..ny - 1 {
270            factxp = (kdint[aij(i + 1, j)] + kdint[aij(i, j)]) / 2.0 * (dt / 2.0) / (dx * dx);
271            factxm = (kdint[aij(i - 1, j)] + kdint[aij(i, j)]) / 2.0 * (dt / 2.0) / (dx * dx);
272            factyp = (kdint[aij(i, j + 1)] + kdint[aij(i, j)]) / 2.0 * (dt / 2.0) / (dy * dy);
273            factym = (kdint[aij(i, j - 1)] + kdint[aij(i, j)]) / 2.0 * (dt / 2.0) / (dy * dy);
274            diag[j] = 1.0 + factyp + factym;
275            sup[j] = -factyp;
276            inf[j] = -factym;
277            f[j] = zint[aij(i, j)] + factxp * zint[aij(i + 1, j)]
278                - (factxp + factxm) * zint[aij(i, j)]
279                + factxm * zint[aij(i - 1, j)];
280        }
281        /*
282        ! bottom bc
283            if (cbc(1:1).eq.'1') then
284            diag(1)=1.
285            sup(1)=0.
286            f(1)=zint(i,1)
287            else
288            factxp=(kdint(i+1,1)+kdint(i,j))/2.d0*(dt/2.)/dx**2
289            factxm=(kdint(i-1,1)+kdint(i,1))/2.d0*(dt/2.)/dx**2
290            factyp=(kdint(i,2)+kdint(i,1))/2.d0*(dt/2.)/dy**2
291            diag(1)=1.d0+factyp
292            sup(1)=-factyp
293            f(1)=zint(i,1)+factxp*zint(i+1,1)-(factxp+factxm)*zint(i,1)+factxm*zint(i-1,1)
294            endif
295        */
296        if true {
297            diag[0] = 1.0;
298            sup[0] = 0.0;
299            f[0] = zint[aij(i, 0)];
300        } else {
301            // reflective boundary
302            // TODO: Check whether this j was actually supposed to be a 0 in the original
303            // (probably).
304            // factxp = (kdint[aij(i+1, 0)] + kdint[aij(i, j)]) / 2.0 * (dt / 2.0) / (dx *
305            // dx);
306            factxp = (kdint[aij(i + 1, 0)] + kdint[aij(i, 0)]) / 2.0 * (dt / 2.0) / (dx * dx);
307            factxm = (kdint[aij(i - 1, 0)] + kdint[aij(i, 0)]) / 2.0 * (dt / 2.0) / (dx * dx);
308            factyp = (kdint[aij(i, 1)] + kdint[aij(i, 0)]) / 2.0 * (dt / 2.0) / (dy * dy);
309            diag[0] = 1.0 + factyp;
310            sup[0] = -factyp;
311            f[0] = zint[aij(i, 0)] + factxp * zint[aij(i + 1, 0)]
312                - (factxp + factxm) * zint[aij(i, 0)]
313                + factxm * zint[aij(i - 1, 0)];
314        }
315        /*
316        ! top bc
317            if (cbc(3:3).eq.'1') then
318            diag(ny)=1.
319            inf(ny)=0.
320            f(ny)=zint(i,ny)
321            else
322            factxp=(kdint(i+1,ny)+kdint(i,ny))/2.d0*(dt/2.)/dx**2
323            factxm=(kdint(i-1,ny)+kdint(i,ny))/2.d0*(dt/2.)/dx**2
324            factym=(kdint(i,ny-1)+kdint(i,ny))/2.d0*(dt/2.)/dy**2
325            diag(ny)=1.d0+factym
326            inf(ny)=-factym
327            f(ny)=zint(i,ny)+factxp*zint(i+1,ny)-(factxp+factxm)*zint(i,ny)+factxm*zint(i-1,ny)
328            endif
329        */
330        if true {
331            diag[ny - 1] = 1.0;
332            inf[ny - 1] = 0.0;
333            f[ny - 1] = zint[aij(i, ny - 1)];
334        } else {
335            // reflective boundary
336            factxp =
337                (kdint[aij(i + 1, ny - 1)] + kdint[aij(i, ny - 1)]) / 2.0 * (dt / 2.0) / (dx * dx);
338            factxm =
339                (kdint[aij(i - 1, ny - 1)] + kdint[aij(i, ny - 1)]) / 2.0 * (dt / 2.0) / (dx * dx);
340            factym = (kdint[aij(i, ny - 2)] + kdint[aij(i, ny - 1)]) / 2.0 * (dt / 2.0) / (dy * dy);
341            diag[ny - 1] = 1.0 + factym;
342            inf[ny - 1] = -factym;
343            f[ny - 1] = zint[aij(i, ny - 1)] + factxp * zint[aij(i + 1, ny - 1)]
344                - (factxp + factxm) * zint[aij(i, ny - 1)]
345                + factxm * zint[aij(i - 1, ny - 1)];
346        }
347        /*
348          call tridag (inf,diag,sup,f,res,ny)
349            do j=1,ny
350            zintp(i,j)=res(j)
351            enddo
352        */
353        tridag(&inf, &diag, &sup, &f, &mut res, ny);
354        for j in 0..ny {
355            zintp[aij(i, j)] = res[j];
356        }
357        /*
358          enddo
359          deallocate (f,diag,sup,inf,res)
360        */
361    }
362    /*
363      ! stores result in 1D array
364
365      do j=1,ny
366        do i=1,nx
367          ij=(j-1)*nx+i
368          etot(ij)=etot(ij)+h(ij)-zintp(i,j)
369          erate(ij)=erate(ij)+(h(ij)-zintp(i,j))/dt
370          h(ij)=zintp(i,j)
371        enddo
372      enddo
373
374      b=min(h,b)
375    */
376    for j in 0..ny {
377        for i in 0..nx {
378            ij = aij(i, j);
379            // FIXME: Track total erosion and erosion rate.
380            h[ij] = zintp[ij] as Alt;
381        }
382    }
383
384    b.par_iter_mut().zip(h).for_each(|(b, h)| {
385        *b = h.min(*b);
386    });
387    /*
388      deallocate (zint,kdint,zintp)
389
390      return
391
392    end subroutine Diffusion
393    */
394}
395
396/*
397!----------
398
399! subroutine to solve a tri-diagonal system of equations (from Numerical Recipes)
400
401      SUBROUTINE tridag(a,b,c,r,u,n)
402
403      implicit none
404
405      INTEGER n
406      double precision a(n),b(n),c(n),r(n),u(n)
407*/
408pub fn tridag(a: &[f64], b: &[f64], c: &[f64], r: &[f64], u: &mut [f64], n: usize) {
409    /*
410          INTEGER j
411          double precision bet
412          double precision,dimension(:),allocatable::gam
413
414          allocate (gam(n))
415
416          if(b(1).eq.0.d0) stop 'in tridag'
417    */
418    let mut bet: f64;
419    let mut gam: Vec<f64>;
420
421    gam = vec![Default::default(); n];
422
423    assert_ne!(b[0], 0.0);
424    /*
425
426    ! first pass
427
428    bet=b(1)
429    u(1)=r(1)/bet
430    do 11 j=2,n
431      gam(j)=c(j-1)/bet
432      bet=b(j)-a(j)*gam(j)
433      if(bet.eq.0.) then
434        print*,'tridag failed'
435        stop
436      endif
437      u(j)=(r(j)-a(j)*u(j-1))/bet
438      11    continue
439    */
440    bet = b[0];
441    u[0] = r[0] / bet;
442    for j in 1..n {
443        gam[j] = c[j - 1] / bet;
444        bet = b[j] - a[j] * gam[j];
445        assert_ne!(bet, 0.0);
446        // Round 0: u[0] = r[0] / b[0]
447        //               = r'[0] / b'[0]
448        // Round j: u[j] = (r[j] - a[j] * u'[j - 1]) / b'[j]
449        //               = (r[j] - a[j] * r'[j - 1] / b'[j - 1]) / b'[j]
450        //               = (r[j] - (a[j] / b'[j - 1]) * r'[j - 1]) / b'[j]
451        //               = (r[j] - w[j] * r'[j - 1]) / b'[j]
452        //               = r'[j] / b'[j]
453        u[j] = (r[j] - a[j] * u[j - 1]) / bet;
454    }
455    /*
456      ! second pass
457
458      do 12 j=n-1,1,-1
459        u(j)=u(j)-gam(j+1)*u(j+1)
460        12    continue
461    */
462    for j in (0..n - 1).rev() {
463        u[j] -= gam[j + 1] * u[j + 1];
464    }
465    /*
466        deallocate (gam)
467
468        return
469
470        END
471    */
472}