From 26b8385635139b2255ccb526f2e320fb15dd9c18 Mon Sep 17 00:00:00 2001 From: Corentin Carton de Wiart Date: Tue, 16 Jul 2024 14:55:09 +0100 Subject: [PATCH] minor interface change for mct --- src/lisflood/hydrological_modules/mct.py | 100 +++++++++++++++++++ src/lisflood/hydrological_modules/routing.py | 74 +++----------- 2 files changed, 114 insertions(+), 60 deletions(-) diff --git a/src/lisflood/hydrological_modules/mct.py b/src/lisflood/hydrological_modules/mct.py index 004c0fd..fa6715c 100644 --- a/src/lisflood/hydrological_modules/mct.py +++ b/src/lisflood/hydrological_modules/mct.py @@ -1,6 +1,106 @@ import numpy as np +import pandas as pd from numba import njit, prange +from .kinematic_wave_parallel import rebuildFlowMatrix, decodeFlowMatrix, streamLookups, topoDistFromSea + + +class MCTWave: + """Build pixels loop for MCT channels""" + + def __init__( + self, + compressed_encoded_ldd, + land_mask, + ChanLength, + ChanGrad, + ChanBottomWidth, + ChanManMCT, + ChanSdXdY, + dt, # computation time step for channel [s] + river_router, + ): + """""" + # Process flow direction matrix: downstream and upstream lookups, and routing orders + flow_dir = decodeFlowMatrix(rebuildFlowMatrix(compressed_encoded_ldd, land_mask)) + self.downstream_lookup, self.upstream_lookup = streamLookups(flow_dir, land_mask) + self.num_upstream_pixels = (self.upstream_lookup != -1).sum(1).astype(int) # astype for cython import in windows (to avoid 'long long' buffer dtype mismatch) + # Routing order: decompose domain into batches; within each batch, pixels can be routed in parallel + self._setMCTRoutingOrders() + + self.ChanLength = ChanLength + self.ChanGrad = ChanGrad + self.ChanBottomWidth = ChanBottomWidth + self.ChanManMCT = ChanManMCT + self.ChanSdXdY = ChanSdXdY + self.dt = dt + self.river_router = river_router + + # create mapping from global domain pixels index to MCT pixels index + idpix_kin = range(len(self.ChanLength)) + self.mapping_mct = self.compress_mct(idpix_kin) + + def _setMCTRoutingOrders(self): + """Compute the MCT wave routing order. Pixels are grouped in sets with the same order. + Pixels in the same set are independent and can be routed in parallel. Sets must be processed in series, starting from order 0. + Pixels are ordered topologically starting from the outlets, as in: + Liu et al. (2014), A layered approach to parallel computing for spatially distributed hydrological modeling, + Environmental Modelling & Software 51, 221-227. + Order MAX is given to pixels with no downstream relations (outlets); order MAX-1 is given to + pixels whose downstream pixels are all of order MAX; and so on.""" + + ocean_topo_distance = topoDistFromSea(self.downstream_lookup, self.upstream_lookup) + routing_order = ocean_topo_distance.max() - ocean_topo_distance + self.pixels_ordered = pd.DataFrame({"pixels": np.arange(routing_order.size), "order": routing_order}) + try: + self.pixels_ordered = self.pixels_ordered.sort_values(["order", "pixels"]).set_index("order").squeeze() + except: # FOR COMPATIBILITY WITH OLDER PANDAS VERSIONS + self.pixels_ordered = self.pixels_ordered.sort(["order", "pixels"]).set_index("order").squeeze() + # Ensure output is always a Series, even if only one element + if not isinstance(self.pixels_ordered, pd.Series): + self.pixels_ordered = pd.Series(self.pixels_ordered) + + order_counts = self.pixels_ordered.groupby(self.pixels_ordered.index).count() + stop = order_counts.cumsum() + self.order_start_stop = np.column_stack((np.append(0, stop[:-1]), stop)).astype(int) # astype for cython import in windows (see above) + self.pixels_ordered = self.pixels_ordered.values.astype(int) # astype for cython import in windows (see above) + + def routing( + self, + ChanQ_0, # q10 + ChanM3_0, # V00 + SideflowChanMCT, + # THESE ARE USED AS INPUTS AND OUTPUTS + ChanQ, # q01 as input, q11 as output + ChanQAvgDt, # q0m as input, q1m as output + PrevCm0, + PrevDm0, + ChanM3, # V11 as output + ): + + mct_routing( + self.ChanLength, + self.ChanGrad, + self.ChanBottomWidth, + self.ChanManMCT, + self.ChanSdXdY, + self.dt, + self.order_start_stop, + self.pixels_ordered, + self.river_router.upstream_lookup, + self.river_router.num_upstream_pixels, + self.mapping_mct, + ChanQ_0, # q10 + ChanM3_0, # V00 + SideflowChanMCT, + # THESE ARE USED AS INPUTS AND OUTPUTS + ChanQ, # q01 as input, q11 as output + ChanQAvgDt, # q0m as input, q1m as output + PrevCm0, + PrevDm0, + ChanM3, # V11 as output + ) + @njit(parallel=True, fastmath=False, cache=True) def mct_routing( diff --git a/src/lisflood/hydrological_modules/routing.py b/src/lisflood/hydrological_modules/routing.py index f0b5458..2b0d586 100644 --- a/src/lisflood/hydrological_modules/routing.py +++ b/src/lisflood/hydrological_modules/routing.py @@ -27,7 +27,7 @@ from .inflow import inflow from .transmission import transmission from .kinematic_wave_parallel import kinematicWave, kwpt -from .mct import mct_routing +from .mct import MCTWave from ..global_modules.settings import LisSettings, MaskInfo from ..global_modules.add1 import loadmap, loadmap_base, compressArray, decompress @@ -573,11 +573,19 @@ def initialMCT(self): # ************************************************************ # ***** INITIALISE MUSKINGUM-CUNGE-TODINI WAVE ROUTER ******** # ************************************************************ - self.mct_river_router = mctWave(self.compress_mct(compressArray(self.var.LddMCT)), self.var.mctmask) + mct_ldd = self.compress_mct(compressArray(self.var.LddMCT)) + self.mct_river_router = MCTWave( + mct_ldd, + self.var.mctmask, + self.var.ChanLength, + self.var.ChanGrad, + self.var.ChanBottomWidth, + self.var.ChanManMCT, + self.var.ChanSdXdY, + self.var.DtRouting, + self.river_router + ) - # create mapping from global domain pixels index to MCT pixels index - idpix_kin = range(len(self.var.ChanQ)) - self.mapping_mct = self.compress_mct(idpix_kin) # -------------------------------------------------------------------------- # -------------------------------------------------------------------------- @@ -596,7 +604,6 @@ def dynamic(self, NoRoutingExecuted): self.inflow_module.dynamic_inloop(NoRoutingExecuted) self.transmission_module.dynamic_inloop() - if not(option['dynamicWave']): # ************************************************************ @@ -753,19 +760,7 @@ def dynamic(self, NoRoutingExecuted): self.var.ChanQAvgDt = ChanQAvgDt # Update current state at MCT points - # ChanQ, ChanQAvgDt, ChanM3, PrevCm0, PrevDm0 = mct_routing( - mct_routing( - self.var.ChanLength, - self.var.ChanGrad, - self.var.ChanBottomWidth, - self.var.ChanManMCT, - self.var.ChanSdXdY, - self.var.DtRouting, - self.mct_river_router.order_start_stop, - self.mct_river_router.pixels_ordered, - self.river_router.upstream_lookup, - self.river_router.num_upstream_pixels, - self.mapping_mct, + self.mct_river_router.routing( ChanQ_0, # q10 ChanM3_0, # V00 SideflowChanMCT, @@ -1076,44 +1071,3 @@ def compress_mct(self, var): y = x.compressed() return y - - -from .kinematic_wave_parallel import rebuildFlowMatrix, decodeFlowMatrix, streamLookups, topoDistFromSea -import pandas as pd -class mctWave: - """Build pixels loop for MCT channels""" - - def __init__(self, compressed_encoded_ldd, land_mask): - """""" - # Process flow direction matrix: downstream and upstream lookups, and routing orders - flow_dir = decodeFlowMatrix(rebuildFlowMatrix(compressed_encoded_ldd, land_mask)) - self.downstream_lookup, self.upstream_lookup = streamLookups(flow_dir, land_mask) - self.num_upstream_pixels = (self.upstream_lookup != -1).sum(1).astype(int) # astype for cython import in windows (to avoid 'long long' buffer dtype mismatch) - # Routing order: decompose domain into batches; within each batch, pixels can be routed in parallel - self._setMCTRoutingOrders() - - def _setMCTRoutingOrders(self): - """Compute the MCT wave routing order. Pixels are grouped in sets with the same order. - Pixels in the same set are independent and can be routed in parallel. Sets must be processed in series, starting from order 0. - Pixels are ordered topologically starting from the outlets, as in: - Liu et al. (2014), A layered approach to parallel computing for spatially distributed hydrological modeling, - Environmental Modelling & Software 51, 221-227. - Order MAX is given to pixels with no downstream relations (outlets); order MAX-1 is given to - pixels whose downstream pixels are all of order MAX; and so on.""" - - ocean_topo_distance = topoDistFromSea(self.downstream_lookup, self.upstream_lookup) - routing_order = ocean_topo_distance.max() - ocean_topo_distance - self.pixels_ordered = pd.DataFrame({"pixels": np.arange(routing_order.size), "order": routing_order}) - try: - self.pixels_ordered = self.pixels_ordered.sort_values(["order", "pixels"]).set_index("order").squeeze() - except: # FOR COMPATIBILITY WITH OLDER PANDAS VERSIONS - self.pixels_ordered = self.pixels_ordered.sort(["order", "pixels"]).set_index("order").squeeze() - # Ensure output is always a Series, even if only one element - if not isinstance(self.pixels_ordered, pd.Series): - self.pixels_ordered = pd.Series(self.pixels_ordered) - - order_counts = self.pixels_ordered.groupby(self.pixels_ordered.index).count() - stop = order_counts.cumsum() - self.order_start_stop = np.column_stack((np.append(0, stop[:-1]), stop)).astype(int) # astype for cython import in windows (see above) - self.pixels_ordered = self.pixels_ordered.values.astype(int) # astype for cython import in windows (see above) -