Coverage for C:\src\imod-python\imod\mf6\auxiliary_variables.py: 96%

1 from imod.mf6.interfaces.ipackage import IPackage

2

3

4 def get_variable_names(package: IPackage):

5 auxiliaries = _get_auxiliary_data_variable_names_mapping(

6 package

7 ) # returns something like {"concentration": "species"}

8

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 []

22

23

24 def _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

29

30

31 def 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 """

39

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 )

50

51

52 def 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)