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

27 statements  

« prev     ^ index     » next       coverage.py v7.5.1, created at 2024-05-08 14:15 +0200

1from imod.mf6.interfaces.ipackage import IPackage 

2 

3 

4def 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 

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 

29 

30 

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

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 

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)