Coverage for cvpack/base_custom_function.py: 100%

1""" 

2.. class:: BaseCustomFunction 

3 :platform: Linux, MacOS, Windows 

4 :synopsis: Abstract class for collective variables defined by a custom function 

5 

6.. classauthor:: Charlles Abreu <craabreu@gmail.com> 

7 

8""" 

9 

10from __future__ import annotations 

11 

12import typing as t 

13 

14from .collective_variable import CollectiveVariable 

15from .units import ScalarQuantity, VectorQuantity, value_in_md_units 

16 

17 

18class BaseCustomFunction(CollectiveVariable): 

19 """ 

20 Abstract class for collective variables defined by a custom function. 

21 """ 

22 

23 def _extractParameters( 

24 self, 

25 size: int, 

26 **parameters: ScalarQuantity | VectorQuantity, 

27 ) -> tuple[dict[str, ScalarQuantity], dict[str, list[ScalarQuantity]]]: 

28 global_parameters = {} 

29 perbond_parameters = {} 

30 for name, value in parameters.items(): 

31 data = value_in_md_units(value) 

32 if isinstance(data, t.Iterable): 

33 data = list(data) 

34 if len(data) != size: 

35 raise ValueError( 

36 f"Parameter {name} has {len(data)} values, but there are " 

37 f"{size} groups" 

38 ) 

39 perbond_parameters[name] = data 

40 else: 

41 global_parameters[name] = data 

42 return global_parameters, perbond_parameters 

43 

44 def _addParameters( 

45 self, 

46 overalls: dict[str, ScalarQuantity], 

47 perbonds: dict[str, list[ScalarQuantity]], 

48 groups: list[tuple[int, ...]], 

49 pbc: bool, 

50 ) -> None: 

51 

52 for name, value in overalls.items(): 

53 self.addGlobalParameter(name, value) 

54 for name in perbonds: 

55 self.addPerBondParameter(name) 

56 for group, *values in zip(groups, *perbonds.values()): 

57 self.addBond(group, values) 

58 self.setUsesPeriodicBoundaryConditions(pbc)