diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md index c34eb45..9e781e6 100644 --- a/CONTRIBUTORS.md +++ b/CONTRIBUTORS.md @@ -3,3 +3,4 @@ | GitHub user | Real Name | Affiliation | Date | | ----------- | --------- | ----------- | ---- | | james-bruten-mo | James Bruten | Met Office | 2025-12-09 | +| joewallwork | Joseph Wallwork | University of Cambridge | 2026-01-23 | diff --git a/src/science/core/chemistry/asad/asad_mod.F90 b/src/science/core/chemistry/asad/asad_mod.F90 index 626f861..5637920 100644 --- a/src/science/core/chemistry/asad/asad_mod.F90 +++ b/src/science/core/chemistry/asad/asad_mod.F90 @@ -166,7 +166,7 @@ MODULE asad_mod REAL, ALLOCATABLE :: frph(:) ! fractional product array (het) REAL, ALLOCATABLE :: frpx(:) ! fractional product array (total) ! sparse algebra -INTEGER, ALLOCATABLE :: nonzero_map_unordered(:,:) ! Map of nonzero entries +INTEGER, ALLOCATABLE :: permuted_nonzero_map(:,:) ! permuted map of nonzeros INTEGER, ALLOCATABLE :: modified_map(:,:) ! modified map (after decomposition) INTEGER, ALLOCATABLE :: nonzero_map(:,:) ! Map of nonzero entries, before reordering @@ -544,8 +544,8 @@ SUBROUTINE asad_mod_init(n_points) njcoth(:,:) = 0 IF (method == int_method_NR) THEN - IF (.NOT. ALLOCATED(nonzero_map_unordered)) & - ALLOCATE(nonzero_map_unordered(jpcspf, jpcspf)) + IF (.NOT. ALLOCATED(permuted_nonzero_map)) & + ALLOCATE(permuted_nonzero_map(jpcspf, jpcspf)) IF (.NOT. ALLOCATED(modified_map)) ALLOCATE(modified_map(jpcspf, jpcspf)) IF (.NOT. ALLOCATED(nonzero_map)) ALLOCATE(nonzero_map(jpcspf, jpcspf)) IF (.NOT. ALLOCATED(reorder)) ALLOCATE(reorder(jpcspf)) diff --git a/src/science/core/chemistry/asad/asad_sparse_vars.F90 b/src/science/core/chemistry/asad/asad_sparse_vars.F90 index 8bc69df..a48a9f7 100644 --- a/src/science/core/chemistry/asad/asad_sparse_vars.F90 +++ b/src/science/core/chemistry/asad/asad_sparse_vars.F90 @@ -81,20 +81,21 @@ SUBROUTINE setup_spfuljac() ! ! 1. Create nonzero_map(1:jpcspf,1:jpcspf) and base_tracer(1:spfjsize_max) ! -! nonzero_map(1:jpcspf,1:jpcspf) is an integer indexing array such that -! spfj(nonzero_map(i,j)) = A(i,j) where A is dense Jacobian matrix and -! spfj is the compressed (sparse) storage of the dense Jacobian matrix. +! nonzero_map(1:jpcspf,1:jpcspf) is an integer indexing array such that, for +! a given species index jl, spfj(jl,nonzero_map(i,j)) = A(i,j) where A is +! the dense Jacobian matrix for species jl and spfj is the compressed +! (sparse) storage of the dense Jacobian matrix. ! -! 2. Create nonzero_map_unordered(1:jpcspf,1:jpcspf) +! 2. Create permuted_nonzero_map(1:jpcspf,1:jpcspf) ! -! nonzero_map_unordered(1:jpcspf,1:jpcspf) is an integer indexing array such -! that spfj(nonzero_map_unordered(i,j)) = (PAP')(i,j) where (PAP')(i,j) is +! permuted_nonzero_map(1:jpcspf,1:jpcspf) is an integer indexing array such +! that spfj(jl,permuted_nonzero_map(i,j)) = (PAP')(i,j) where (PAP')(i,j) is ! the (i,j) entry of the matrix matmul(P,matmul(A,transpose(P))) where P is ! the permutation matrix and A is the dense Jacobian. ! ! 3. Check number of nonzero entries in LU factorization of the dense Jacobian ! -! Gaussian-elimination is used to solve a linear equation in splinslv2 and +! Gaussian elimination is used to solve a linear equation in splinslv2 and ! the array spfj is used to hold the LU factorization. This section checks ! that when this LU factorization is done the spfj array is sufficiently ! large enough to hold the nonzero entries of the LU factorization. @@ -106,7 +107,7 @@ SUBROUTINE setup_spfuljac() USE asad_mod, ONLY: specf, frpx, jpcspf, jpfrpx, jpmsp, jpspec, & madvtr, modified_map, ndepd, ndepw, nfrpx, njcoth, nltrf, & - nmsjac, nmzjac, nonzero_map, nonzero_map_unordered, & + nmsjac, nmzjac, nonzero_map, permuted_nonzero_map, & npdfr, nsjac1, nstst, ntabpd, ntrf, ntro3, nzjac1, & reorder, spfjsize_max, maxterms, maxfterms, & nposterms, nnegterms, nfracterms, posterms, negterms, & @@ -183,10 +184,9 @@ SUBROUTINE setup_spfuljac() IF (ij /= 0) map(ij,itrcr) = 1 END DO ! - IF (npdfr(irj,1) /= 0) THEN - i1 = npdfr(irj,1) - i2 = npdfr(irj,2) - DO jn = i1, i2 + i1 = npdfr(irj,1) + IF (i1 /= 0) THEN + DO jn = i1, npdfr(irj,2) isx = ntabpd(jn,1) map(isx,itrcr) = 1 END DO @@ -241,15 +241,15 @@ SUBROUTINE setup_spfuljac() END DO base_tracer(:) = 0 -DO i = 1, jpcspf - DO j = 1, jpcspf - i1 = nonzero_map(j,i) - IF (i1 > 0) base_tracer(i1) = i +DO j = 1, jpcspf + DO i = 1, jpcspf + i1 = nonzero_map(i,j) + IF (i1 > 0) base_tracer(i1) = j END DO END DO ! ------------------------------------------------------------------------------ -! Section 2: Create nonzero_map_unordered (index array for PAP') +! Section 2: Create permuted_nonzero_map (index array for PAP') ! ------------------------------------------------------------------------------ ! Reorder species by their reactivity to minimize fill-in @@ -280,19 +280,19 @@ SUBROUTINE setup_spfuljac() END IF END DO -! reorganize pointer variable to account for varying fill-in +! Determine permutation matrix P IF (.NOT. ALLOCATED(permute)) ALLOCATE(permute(jpcspf,jpcspf)) - permute(:,:) = 0 DO i = 1, jpcspf permute(i,reorder(i)) = 1 END DO -! Calculate the index array nonzero_map_unordered such that -! spfj(nonzero_map_unordered(i,j)) = P*A*P'(i,j) -! where P is the permutation matrix and A is the dense Jacobian matrix -! with spfj the compressed storage (sparse) array for matrix A. -nonzero_map_unordered = MATMUL(MATMUL(permute, nonzero_map), TRANSPOSE(permute)) +! Calculate the index array permuted_nonzero_map such that +! spfj(jl,permuted_nonzero_map(i,j)) = P*A*P'(i,j) +! for each chemical species jl where P is the permutation matrix and A is the +! dense Jacobian matrix with spfj the compressed storage (sparse) array for +! matrix A. +permuted_nonzero_map = MATMUL(MATMUL(permute, nonzero_map), TRANSPOSE(permute)) IF (ALLOCATED(permute)) DEALLOCATE(permute) IF (ALLOCATED(map)) DEALLOCATE(map) @@ -304,7 +304,7 @@ SUBROUTINE setup_spfuljac() ! Calculate the number of nonzero matrix elements in the LU factorization of the ! array PAP' and check that it is less than spfjsize_max. total1 = total -modified_map(:,:) = nonzero_map_unordered(:,:) +modified_map(:,:) = permuted_nonzero_map(:,:) DO kr = 1, jpcspf DO i = kr+1, jpcspf ikr = modified_map(i,kr) @@ -313,8 +313,7 @@ SUBROUTINE setup_spfuljac() krj = modified_map(kr,j) IF (krj > 0) THEN ij = modified_map(i,j) - ! Distinguish whether matrix element is zero or not. If not, proceed - ! as in dense case. If it is, create new matrix element. + ! If the matrix entry is zero then create new matrix element. IF (ij <= 0) THEN total1 = total1 + 1 modified_map(i,j) = total1 @@ -561,9 +560,9 @@ SUBROUTINE spfuljac(n_points, cdt, min_pivot, nonzero_map, spfj) DO js = 1, jpspec ifamd = moffam(js) itrd = madvtr(js) - ityped = ctype(js) ! IF (ifamd /= 0) THEN + ityped = ctype(js) p = nonzero_map(ifamd,ifamd) DO jl = 1, n_points IF ((ityped == jpfm) .OR. ((ityped == jpif) .AND. linfam(jl,itrd))) THEN @@ -594,7 +593,7 @@ END SUBROUTINE spfuljac !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! SUBROUTINE splinslv2(n_points, bb, xx, min_pivot, max_val, & - nonzero_map_unordered, modified_map, spfj) + permuted_nonzero_map, modified_map, spfj) USE asad_mod, ONLY: jpcspf, spfjsize_max, total USE parkind1, ONLY: jprb, jpim @@ -629,12 +628,11 @@ SUBROUTINE splinslv2(n_points, bb, xx, min_pivot, max_val, & REAL, INTENT(OUT) :: xx(n_points,jpcspf) REAL, INTENT(IN) :: min_pivot REAL, INTENT(IN) :: max_val -INTEGER, INTENT(IN) :: nonzero_map_unordered(jpcspf,jpcspf) +INTEGER, INTENT(IN) :: permuted_nonzero_map(jpcspf,jpcspf) INTEGER, INTENT(IN OUT) :: modified_map(jpcspf,jpcspf) REAL, INTENT(IN OUT) :: spfj(1:n_points,1:spfjsize_max) ! Local variables -INTEGER :: total1 ! total number of nonzero entries in Jacobian INTEGER :: kr INTEGER :: jl INTEGER :: i @@ -643,11 +641,9 @@ SUBROUTINE splinslv2(n_points, bb, xx, min_pivot, max_val, & INTEGER :: krj INTEGER :: ij -REAL :: bb1(n_points,jpcspf) -REAL :: xx1(n_points,jpcspf) - +REAL :: multiplier(n_points) REAL :: pivot(n_points) -REAL :: kfact(n_points) +REAL :: lower(n_points) INTEGER(KIND=jpim), PARAMETER :: zhook_in = 0 INTEGER(KIND=jpim), PARAMETER :: zhook_out = 1 @@ -657,6 +653,10 @@ SUBROUTINE splinslv2(n_points, bb, xx, min_pivot, max_val, & IF (lhook) CALL dr_hook(ModuleName//':'//RoutineName,zhook_in,zhook_handle) +! NOTE: splinslv2 is always preceded by spfuljac, which computes spfj to be a +! sparse (compressed) representation of the Jacobian matrix for each +! species. + ! Filter sparse Jacobian #if defined(IBM_XL_FORTRAN) ! Version optimised for IBM by using the fsel IBM-only intrinsic @@ -676,35 +676,46 @@ SUBROUTINE splinslv2(n_points, bb, xx, min_pivot, max_val, & ! Section 1: Determine L U factors such that L U = P A P' -! The L U factors are overwritten onto the original sparse Jacobian array -total1 = total -modified_map(:,:) = nonzero_map_unordered(:,:) +! +! The L U factors are overwritten onto the original sparse Jacobian array. That +! is, for each species index jl, the upper triangular part of the matrix +! represented by spfj(jl,:) holds the U factor and the strictly lower triangular +! part holds the L factor below its unit diagonal. + +! Loop over rows for pivoting DO kr = 1, jpcspf - pivot(:) = spfj(:,modified_map(kr,kr)) - WHERE (ABS(pivot) > min_pivot) - pivot(:) = 1.0 / pivot + + ! The multipliers are the diagonal entries for each species + multiplier(:) = spfj(:,modified_map(kr,kr)) + + ! Determine the pivot for each species using a threshold + WHERE (ABS(multiplier) > min_pivot) + pivot(:) = 1.0 / multiplier ELSE WHERE pivot(:) = max_val END WHERE - ! PIVOT = 1./spfj(:,modified_map(kr,kr)) + + ! Loop over non-zero entries in the lower-diagonal part of column kr DO i = kr+1, jpcspf ikr = modified_map(i,kr) IF (ikr > 0) THEN - kfact = spfj(:,ikr)*pivot - spfj(:,ikr) = kfact + + ! Compute contribution to the lower diagonal part and store it in spfj + lower(:) = spfj(:,ikr)*pivot + spfj(:,ikr) = lower + + ! Loop over non-zero entries in the upper-diagonal part of row kr DO j = kr+1, jpcspf krj = modified_map(kr,j) IF (krj > 0) THEN + + ! Compute the contribution towards the upper triangular part and store + ! it in the upper-diagonal part of spfj + ! NOTE: We can rely on the entry being nonzero because the sparsity + ! pattern was pre-computed. ij = modified_map(i,j) - ! Distinguish whether matrix element is zero or not. If not, proceed - ! as in dense case. If it is, create new matrix element. - IF (ij > 0) THEN - spfj(:,ij) = spfj(:,ij) - kfact*spfj(:,krj) - ELSE - total1 = total1 + 1 - modified_map(i,j) = total1 - spfj(:,total1) = -kfact*spfj(:,krj) - END IF + spfj(:,ij) = spfj(:,ij) - lower*spfj(:,krj) + END IF END DO END IF @@ -714,14 +725,14 @@ SUBROUTINE splinslv2(n_points, bb, xx, min_pivot, max_val, & ! Filter sparse Jacobian #if defined(IBM_XL_FORTRAN) ! Version optimised for IBM by using the fsel IBM-only intrinsic -DO j = 1, total1 +DO j = 1, total DO jl = 1, n_points tmp = fsel(spfj(jl,j)+max_val, spfj(jl,j), -max_val) spfj(jl,j) = fsel(tmp-max_val, max_val, tmp) END DO END DO #else -DO j = 1, total1 +DO j = 1, total DO jl = 1, n_points spfj(jl,j) = MIN(MAX(spfj(jl,j), -max_val), max_val) END DO diff --git a/src/science/core/chemistry/asad/asad_spimpmjp.F90 b/src/science/core/chemistry/asad/asad_spimpmjp.F90 index 6d28029..86e46ff 100644 --- a/src/science/core/chemistry/asad/asad_spimpmjp.F90 +++ b/src/science/core/chemistry/asad/asad_spimpmjp.F90 @@ -257,7 +257,7 @@ SUBROUTINE asad_spimpmjp(exit_code, ix, jy, nlev, n_points, location, & USE asad_mod, ONLY: ptol, peps, cdt, f, fdot, nitnr, nstst, y, & fj, nonzero_map, ltrig, jpcspf, spfj, & - modified_map, nonzero_map_unordered + modified_map, permuted_nonzero_map USE asad_sparse_vars, ONLY: setup_spfuljac, spfuljac, spresolv2, splinslv2 USE ukca_config_specification_mod, ONLY: ukca_config USE yomhook, ONLY: lhook, dr_hook @@ -470,8 +470,8 @@ SUBROUTINE asad_spimpmjp(exit_code, ix, jy, nlev, n_points, location, & END DO END IF - CALL splinslv2(n_points,G_f,f_incr,f_min,f_max,nonzero_map_unordered, & - modified_map,spfj) + CALL splinslv2(n_points,G_f,f_incr,f_min,f_max,permuted_nonzero_map, & + modified_map,spfj) IF (ltrig .AND. printstatus == PrStatus_Diag) THEN DO jl=1,n_points diff --git a/src/science/core/chemistry/asad/asad_spmjpdriv.F90 b/src/science/core/chemistry/asad/asad_spmjpdriv.F90 index 6197c73..6b10bed 100644 --- a/src/science/core/chemistry/asad/asad_spmjpdriv.F90 +++ b/src/science/core/chemistry/asad/asad_spmjpdriv.F90 @@ -117,6 +117,7 @@ SUBROUTINE asad_spmjpdriv(ix,jy,nlev,n_points) USE errormessagelength_mod, ONLY: errormessagelength +! TODO: This import will become redundant USE asad_diffun_mod, ONLY: asad_diffun USE asad_spimpmjp_mod, ONLY: asad_spimpmjp USE asad_ftoy_mod, ONLY: asad_ftoy