def _load_external_fluxes(self, **keywords):
func = keywords["func"]
flux_structure = keywords["flux_structure"]
ms = self.model_structure
time_agg = self.time_agg
stock_unit = self.stock_unit
fs_data = np.ma.masked_array(data=np.zeros(
(len(time_agg.data) - 1, ms.nr_pools)),
mask=False)
for pool_name, variable_names in flux_structure.items():
nr_layers = ms.get_nr_layers(pool_name)
dz = self.get_dz(pool_name)
fvs_agg = []
for variable_name in variable_names:
# print(pool_name, variable_name, flush=True)
fv_agg = func(mdo=self,
variable_name=variable_name,
nr_layers=nr_layers,
dz=dz,
**keywords)
fvs_agg.append(fv_agg)
fv_pool_agg = sum(fvs_agg)
for ly in range(nr_layers):
pool_nr = ms.get_pool_nr(pool_name, ly)
data = fv_pool_agg.data[:, ly, ...]
fs_data[..., pool_nr] = data
fs = FluxVariable(data=fs_data, unit=self.stock_unit)
return fs
def load_horizontal_fluxes(self, **keywords):
func = keywords["func"]
ms = self.model_structure
time_agg = self.time_agg
stock_unit = self.stock_unit
HFs_data = np.ma.masked_array(
data=np.zeros((len(time_agg.data) - 1, ms.nr_pools, ms.nr_pools)),
mask=False,
)
for pools, variable_names in ms.horizontal_structure.items():
src_pool_name = pools[0]
tar_pool_name = pools[1]
# print(src_pool_name, tar_pool_name, flush=True)
src_nr_layers = ms.get_nr_layers(src_pool_name)
tar_nr_layers = ms.get_nr_layers(tar_pool_name)
assert src_nr_layers == tar_nr_layers
nr_layers = src_nr_layers
src_dz = self.get_dz(src_pool_name)
tar_dz = self.get_dz(tar_pool_name)
assert src_dz.name == tar_dz.name
dz = src_dz
fvs_agg = []
for variable_name in variable_names:
fv_agg = func(mdo=self,
variable_name=variable_name,
nr_layers=nr_layers,
dz=dz,
**keywords)
fvs_agg.append(fv_agg)
fv_agg = sum(fvs_agg)
for ly in range(nr_layers):
src_pool_nr = ms.get_pool_nr(src_pool_name, ly)
tar_pool_nr = ms.get_pool_nr(tar_pool_name, ly)
# print(src_pool_name, tar_pool_name, ly, src_pool_nr, tar_pool_nr, flush=True)
data = fv_agg.data[:, ly, ...]
HFs_data[..., tar_pool_nr, src_pool_nr] = data
HFs = FluxVariable(name="horizontal fluxes",
data=HFs_data,
unit=self.stock_unit)
return HFs
def test_aggregateInTime(self):
## one-dimensional
masked_data = np.ma.masked_array(
data=np.arange(10), mask=[0, 0, 1, 1, 0, 1, 1, 0, 0, 0], fill_value=-1
)
unit = "kg"
fv = FluxVariable(data=masked_data, unit=unit)
nstep = 1
fv_agg = fv.aggregateInTime(nstep)
res2 = [0, 1, -1, -1, 4, -1, -1, 7, 8, 9]
self.assertTrue(np.all(fv_agg.data.filled() == res2))
nstep = 2
fv_agg = fv.aggregateInTime(nstep)
res2 = np.array([1, -1, -1, -1, 17])
self.assertTrue(np.all(fv_agg.data.filled() == res2))
nstep = 3
fv_agg = fv.aggregateInTime(nstep)
res2 = np.array([-1, -1, -1, 9])
self.assertTrue(np.all(fv_agg.data.filled() == res2))
## multi-dimensional
masked_data = np.ma.masked_array(
data=np.arange(20).reshape((10, 2)),
mask=[
[0, 0],
[0, 1],
[1, 0],
[0, 0],
[0, 0],
[0, 0],
[1, 0],
[0, 0],
[0, 0],
[0, 1],
],
fill_value=-1,
)
unit = "kg"
fv = FluxVariable(data=masked_data, unit=unit)
nstep = 3
fv_agg = fv.aggregateInTime(nstep)
res2 = [[-1, -1], [24, 27], [-1, 45], [18, -1]]
self.assertTrue(np.all(fv_agg.data.filled() == res2))
def load_vertical_fluxes(self, **keywords):
func = keywords["func"]
ms = self.model_structure
time_agg = self.time_agg
stock_unit = self.stock_unit
VFs_data = np.ma.masked_array(
data=np.zeros((len(time_agg.data) - 1, ms.nr_pools, ms.nr_pools)),
mask=False,
)
runoffs_up_data = np.ma.masked_array(data=np.zeros(
(len(time_agg.data) - 1, ms.nr_pools)),
mask=False)
runoffs_down_data = np.ma.masked_array(data=np.zeros(
(len(time_agg.data) - 1, ms.nr_pools)),
mask=False)
src_ly_shift = {
"to_below": 0,
"from_below": 1,
"to_above": 0,
"from_above": -1
}
tar_ly_shift = {
"to_below": 1,
"from_below": 0,
"to_above": -1,
"from_above": 0
}
for pool_name, flux_structure in ms.vertical_structure.items():
nr_layers = ms.get_nr_layers(pool_name)
for flux_type in src_ly_shift.keys():
fvs_agg = []
variable_names = flux_structure.get(flux_type, None)
for variable_name in variable_names:
# print(pool_name, variable_name, flush=True)
fv_agg = func(mdo=self,
variable_name=variable_name,
nr_layers=nr_layers,
**keywords)
fvs_agg.append(fv_agg)
if fvs_agg != []:
fv_agg = sum(fvs_agg)
for ly in range(nr_layers):
pool_nr = ms.get_pool_nr(pool_name, ly)
src_ly = ly + src_ly_shift[flux_type]
tar_ly = ly + tar_ly_shift[flux_type]
if (src_ly
in range(nr_layers)) and (tar_ly
in range(nr_layers)):
src_pool_nr = ms.get_pool_nr(pool_name, src_ly)
tar_pool_nr = ms.get_pool_nr(pool_name, tar_ly)
data = fv_agg.data[:, ly, ...]
VFs_data[..., tar_pool_nr, src_pool_nr] = data
elif src_ly in range(nr_layers) and (tar_ly == -1):
src_pool_nr = ms.get_pool_nr(pool_name, src_ly)
data = fv_agg.data[:, ly, ...]
runoffs_up_data[..., src_pool_nr] = data
elif src_ly in range(nr_layers) and (tar_ly
== nr_layers):
src_pool_nr = ms.get_pool_nr(pool_name, src_ly)
data = fv_agg.data[:, ly, ...]
runoffs_down_data[..., src_pool_nr] = data
VFs = FluxVariable(name="vertical_fluxes",
data=VFs_data,
unit=self.stock_unit)
runoffs_up = FluxVariable(name="runoffs up",
data=runoffs_up_data,
unit=self.stock_unit)
runoffs_down = FluxVariable(name="runoffs down",
data=runoffs_down_data,
unit=self.stock_unit)
return VFs, runoffs_up, runoffs_down
def load_horizontal_fluxes(self, **keywords):
func = keywords["func"]
keywords["check_units"] = self.check_units
ms = self.model_structure
time_agg = self.time_agg
stock_unit = self.stock_unit
HFs_data = np.ma.masked_array(
data=np.zeros((len(time_agg.data) - 1, ms.nr_pools, ms.nr_pools)),
mask=False,
)
for pools, variable_names in ms.horizontal_structure.items():
src_pool_name = pools[0]
tar_pool_name = pools[1]
# print(src_pool_name, tar_pool_name, flush=True)
src_nr_layers = ms.get_nr_layers(src_pool_name)
tar_nr_layers = ms.get_nr_layers(tar_pool_name)
# assert src_nr_layers == tar_nr_layers
if src_nr_layers == tar_nr_layers:
nr_layers = src_nr_layers
src_dz = self.get_dz(src_pool_name)
tar_dz = self.get_dz(tar_pool_name)
assert src_dz.name == tar_dz.name
dz = src_dz
fvs_agg = []
for variable_name in variable_names:
fv_agg = func(mdo=self,
variable_name=variable_name,
nr_layers=nr_layers,
dz=dz,
**keywords)
fvs_agg.append(fv_agg)
fv_agg = sum(fvs_agg)
for ly in range(nr_layers):
src_pool_nr = ms.get_pool_nr(src_pool_name, ly)
tar_pool_nr = ms.get_pool_nr(tar_pool_name, ly)
# print(src_pool_name, tar_pool_name, ly, src_pool_nr, tar_pool_nr, flush=True)
data = fv_agg.data[:, ly, ...]
HFs_data[..., tar_pool_nr, src_pool_nr] = data
elif src_nr_layers == 1:
tar_dz = self.get_dz(tar_pool_name)
fvs_agg = []
for variable_name in variable_names:
fv_agg = func(mdo=self,
variable_name=variable_name,
nr_layers=tar_nr_layers,
dz=tar_dz,
**keywords)
fvs_agg.append(fv_agg)
fv_agg = sum(fvs_agg)
src_pool_nr = ms.get_pool_nr(src_pool_name, 0)
for ly in range(tar_nr_layers):
tar_pool_nr = ms.get_pool_nr(tar_pool_name, ly)
# print(src_pool_name, tar_pool_name, ly, src_pool_nr, tar_pool_nr, flush=True)
data = fv_agg.data[:, ly, ...]
HFs_data[..., tar_pool_nr, src_pool_nr] = data * 2
else:
s = "layer structure for flux between" + \
src_pool_name + " and " + tar_pool_name + \
" not understood"
raise (ModelDataException(s))
HFs = FluxVariable(name="horizontal fluxes",
data=HFs_data,
unit=self.stock_unit)
return HFs