Skip to content

Commit

Permalink
write fluxes one by on all times
Browse files Browse the repository at this point in the history
  • Loading branch information
Apolline El-Baz committed Apr 12, 2024
1 parent fdcb106 commit b5ab0fb
Show file tree
Hide file tree
Showing 8 changed files with 565 additions and 109 deletions.
11 changes: 11 additions & 0 deletions smash/_constant.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -826,6 +826,17 @@ def get_rr_states_from_structure(structure: str) -> list[str]:
"jobs": False,
"qt": False,
"stats": False,
"ei": False,
"pn": False,
"en": False,
"pr": False,
"perc": False,
"lexc": False,
"prr": False,
"prd": False,
"qr": False,
"qd": False,
"qb": False,
},
"optimize": {
"time_step": "all",
Expand Down
88 changes: 88 additions & 0 deletions smash/core/simulation/_doc.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -523,6 +523,94 @@
Whether to return stats : mean, var, min, max, med on internal fluxes.
""",
),
"pn": (
"""
`bool`, default False
""",
"""
Whether to return internal flux pn.
""",
),
"ei": (
"""
`bool`, default False
""",
"""
Whether to return internal flux ei.
""",
),
"en": (
"""
`bool`, default False
""",
"""
Whether to return internal flux en.
""",
),
"pr": (
"""
`bool`, default False
""",
"""
Whether to return internal flux pr.
""",
),
"perc": (
"""
`bool`, default False
""",
"""
Whether to return internal flux perc.
""",
),
"lexc": (
"""
`bool`, default False
""",
"""
Whether to return internal flux lexc.
""",
),
"prr": (
"""
`bool`, default False
""",
"""
Whether to return internal flux prr.
""",
),
"prd": (
"""
`bool`, default False
""",
"""
Whether to return internal flux prd.
""",
),
"qr": (
"""
`bool`, default False
""",
"""
Whether to return internal flux qr.
""",
),
"qd": (
"""
`bool`, default False
""",
"""
Whether to return internal flux qd.
""",
),
"qb": (
"""
`bool`, default False
""",
"""
Whether to return internal flux qb.
""",
),
"iter_cost": (
"""
`bool`, default False
Expand Down
7 changes: 4 additions & 3 deletions smash/core/simulation/run/run.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -164,9 +164,10 @@ def _forward_run(
return_options["nmts"],
return_options["fkeys"],
)

wrap_returns.stats.fluxes_keys = INTERNAL_FLUXES[model.setup.hydrological_module]
wrap_returns.stats.rr_states_keys = STRUCTURE_RR_STATES[model.setup.structure]

if wrap_returns.stats_flag:
wrap_returns.stats.fluxes_keys = INTERNAL_FLUXES[model.setup.hydrological_module]
wrap_returns.stats.rr_states_keys = STRUCTURE_RR_STATES[model.setup.structure]

# % Map cost_options dict to derived type
_map_dict_to_fortran_derived_type(cost_options, wrap_options.cost)
Expand Down
171 changes: 170 additions & 1 deletion smash/fcore/derived_type/mwd_returns.f90
View file
Original file line number Diff line number Diff line change
Expand Up @@ -103,9 +103,42 @@ module mwd_returns
real(sp), dimension(:, :, :), allocatable :: qt
logical :: qt_flag = .false.

! internal fluxes
! stats target on fluxes and states
type (StatsDT) :: stats
logical :: stats_flag = .false.

real(sp), dimension(:, :, :), allocatable :: ei
logical :: ei_flag = .false.

real(sp), dimension(:, :, :), allocatable :: pn
logical :: pn_flag = .false.

real(sp), dimension(:, :, :), allocatable :: en
logical :: en_flag = .false.

real(sp), dimension(:, :, :), allocatable :: pr
logical :: pr_flag = .false.

real(sp), dimension(:, :, :), allocatable :: perc
logical :: perc_flag = .false.

real(sp), dimension(:, :, :), allocatable :: lexc
logical :: lexc_flag = .false.

real(sp), dimension(:, :, :), allocatable :: prr
logical :: prr_flag = .false.

real(sp), dimension(:, :, :), allocatable :: prd
logical :: prd_flag = .false.

real(sp), dimension(:, :, :), allocatable :: qr
logical :: qr_flag = .false.

real(sp), dimension(:, :, :), allocatable :: qd
logical :: qd_flag = .false.

real(sp), dimension(:, :, :), allocatable :: qb
logical :: qb_flag = .false.

end type ReturnsDT

Expand Down Expand Up @@ -195,6 +228,142 @@ subroutine ReturnsDT_initialise(this, setup, mesh, nmts, keys)
call StatsDT_initialise(this%stats, setup, mesh)

end select

if ((setup%hydrological_module == "gr4") .or. (setup%hydrological_module == "gr5")) then
select case (wkeys(i))
case ("pn")
this%pn_flag = .true.
allocate (this%pn(mesh%nrow, mesh%ncol, setup%ntime_step))

case ("en")
this%en_flag = .true.
allocate (this%en(mesh%nrow, mesh%ncol, setup%ntime_step))

case ("pr")
this%pr_flag = .true.
allocate (this%pr(mesh%nrow, mesh%ncol, setup%ntime_step))

case ("perc")
this%perc_flag = .true.
allocate (this%perc(mesh%nrow, mesh%ncol, setup%ntime_step))

case ("lexc")
this%lexc_flag = .true.
allocate (this%lexc(mesh%nrow, mesh%ncol, setup%ntime_step))

case ("prr")
this%prr_flag = .true.
allocate (this%prr(mesh%nrow, mesh%ncol, setup%ntime_step))

case ("prd")
this%prd_flag = .true.
allocate (this%prd(mesh%nrow, mesh%ncol, setup%ntime_step))

case ("qr")
this%qr_flag = .true.
allocate (this%qr(mesh%nrow, mesh%ncol, setup%ntime_step))

case ("qd")
this%qd_flag = .true.
allocate (this%qd(mesh%nrow, mesh%ncol, setup%ntime_step))
end select
end if

if (setup%hydrological_module == "grd") then
select case (wkeys(i))
case ("ei")
this%ei_flag = .true.
allocate (this%ei(mesh%nrow, mesh%ncol, setup%ntime_step))

case ("pn")
this%pn_flag = .true.
allocate (this%pn(mesh%nrow, mesh%ncol, setup%ntime_step))

case ("en")
this%en_flag = .true.
allocate (this%en(mesh%nrow, mesh%ncol, setup%ntime_step))

case ("pr")
this%pr_flag = .true.
allocate (this%pr(mesh%nrow, mesh%ncol, setup%ntime_step))

case ("perc")
this%perc_flag = .true.
allocate (this%perc(mesh%nrow, mesh%ncol, setup%ntime_step))

case ("prr")
this%prr_flag = .true.
allocate (this%prr(mesh%nrow, mesh%ncol, setup%ntime_step))

case ("qr")
this%qr_flag = .true.
allocate (this%qr(mesh%nrow, mesh%ncol, setup%ntime_step))

end select
end if

if (setup%hydrological_module == "loieau") then
select case (wkeys(i))
case ("ei")
this%ei_flag = .true.
allocate (this%ei(mesh%nrow, mesh%ncol, setup%ntime_step))

case ("pn")
this%pn_flag = .true.
allocate (this%pn(mesh%nrow, mesh%ncol, setup%ntime_step))

case ("en")
this%en_flag = .true.
allocate (this%en(mesh%nrow, mesh%ncol, setup%ntime_step))

case ("pr")
this%pr_flag = .true.
allocate (this%pr(mesh%nrow, mesh%ncol, setup%ntime_step))

case ("perc")
this%perc_flag = .true.
allocate (this%perc(mesh%nrow, mesh%ncol, setup%ntime_step))

case ("prr")
this%prr_flag = .true.
allocate (this%prr(mesh%nrow, mesh%ncol, setup%ntime_step))

case ("prd")
this%prd_flag = .true.
allocate (this%prd(mesh%nrow, mesh%ncol, setup%ntime_step))

case ("qr")
this%qr_flag = .true.
allocate (this%qr(mesh%nrow, mesh%ncol, setup%ntime_step))

case ("qd")
this%qd_flag = .true.
allocate (this%qd(mesh%nrow, mesh%ncol, setup%ntime_step))

end select
end if

if (setup%hydrological_module == "vic3l") then
select case (wkeys(i))
case ("pn")
this%pn_flag = .true.
allocate (this%pn(mesh%nrow, mesh%ncol, setup%ntime_step))

case ("en")
this%en_flag = .true.
allocate (this%en(mesh%nrow, mesh%ncol, setup%ntime_step))

case ("qr")
this%qr_flag = .true.
allocate (this%qr(mesh%nrow, mesh%ncol, setup%ntime_step))

case ("qb")
this%qb_flag = .true.
allocate (this%qb(mesh%nrow, mesh%ncol, setup%ntime_step))

end select
end if

end do

end subroutine ReturnsDT_initialise
Expand Down
Loading

0 comments on commit b5ab0fb

Please sign in to comment.