Coverage for C:\src\imod-python\imod\mf6\auxiliary_variables.py: 96%
27 statements
« prev ^ index » next coverage.py v7.4.4, created at 2024-04-08 13:27 +0200
« prev ^ index » next coverage.py v7.4.4, created at 2024-04-08 13:27 +0200
1from imod.mf6.interfaces.ipackage import IPackage
4def get_variable_names(package: IPackage):
5 auxiliaries = _get_auxiliary_data_variable_names_mapping(
6 package
7 ) # returns something like {"concentration": "species"}
9 # loop over the types of auxiliary variables (for example concentration)
10 for auxvar in auxiliaries.keys():
11 # if "concentration" is a variable of this dataset
12 if auxvar in package.dataset.data_vars:
13 # if our concentration dataset has the species coordinate
14 if auxiliaries[auxvar] in package.dataset[auxvar].coords:
15 # assign the species names list to d
16 return package.dataset[auxiliaries[auxvar]].values.tolist()
17 else:
18 raise ValueError(
19 f"{auxvar} requires a {auxiliaries[auxvar]} coordinate."
20 )
21 return []
24def _get_auxiliary_data_variable_names_mapping(package: IPackage):
25 result = {}
26 if hasattr(package, "_auxiliary_data"):
27 result.update(package._auxiliary_data)
28 return result
31def expand_transient_auxiliary_variables(package: IPackage) -> None:
32 """
33 Splits an auxiliary dataarray of the kind that could potentially be
34 time-dependent (with one or more auxiliary variable dimension) into
35 dataarrays per auxiliary variable dimension. For example a concentration
36 auxiliary variable in a flow package will have a species dimension, and will
37 be split in several dataarrays- one for each species.
38 """
40 if len(package.auxiliary_data_fields) > 0:
41 for aux_var_name, aux_var_dimensions in package.auxiliary_data_fields.items():
42 if aux_var_name in list(package.dataset.keys()):
43 aux_coords = (
44 package.dataset[aux_var_name].coords[aux_var_dimensions].values
45 )
46 for s in aux_coords:
47 package.dataset[s] = package.dataset[aux_var_name].sel(
48 {aux_var_dimensions: s}
49 )
52def remove_expanded_auxiliary_variables_from_dataset(package: IPackage) -> None:
53 """
54 Removes the data arrays created by :meth:expand_transient_auxiliary_variables(...) but does not
55 remove the auxiliary dataarray used as source for :meth:expand_transient_auxiliary_variables(...)
56 """
57 for aux_var_name, aux_var_dimensions in package.auxiliary_data_fields.items():
58 if aux_var_dimensions in package.dataset.coords:
59 for species in package.dataset.coords[aux_var_dimensions].values:
60 if species in list(package.dataset.keys()):
61 package.dataset = package.dataset.drop_vars(species)