Coverage for strongcoca / env.py: 100%

1from os import getenv 

2 

3 

4def strongcoca_getenv(variable) -> str: 

5 """ Get value of environment variable, or default. 

6 

7 Possible environment variables are: 

8 

9 - `STRONGCOCA_MAX_SOLVE_MEM` - Value in units of MiB. In the linear solve operation for the 

10 PolarizabilityCalculator, try to limit memory of intermediate matrices below this value. 

11 - `STRONGCOCA_GPU_BACKEND` - GPU backend for the O(N³) linear-algebra operations in 

12 PolarizabilityCalculator. ``'none'`` (default, CPU only), ``'torch'`` (PyTorch/CUDA), 

13 or ``'cupy'`` (CuPy/CUDA). 

14 - `STRONGCOCA_GPU_PRECISION` - Floating-point precision for GPU operations. 

15 ``'float32'`` (default, ~10× faster on consumer GPUs) or ``'float64'``. 

16 - `STRONGCOCA_GPU_BATCHED_MAX` - Maximum N3 = 3×N_particles for which the 

17 cuBLAS batched solver is used in GPU kernels. Above this threshold the 

18 code loops over the batch dimension to invoke the non-batched cuSOLVER 

19 path, which is faster for large matrices on consumer GPUs. Increase this 

20 value (e.g. to 1024 or higher) on A100/H100/H200 cards where the batched 

21 cuSOLVER path performs well at large N. Default: 32. 

22 

23 Parameters 

24 ---------- 

25 variable 

26 Name of variable, without the ``'STRONGCOCA'`` prefix. 

27 

28 Returns 

29 ------- 

30 Value of variable as string. 

31 """ 

32 

33 defaults = {'MAX_SOLVE_MEM': '80', 

34 'GPU_BACKEND': 'none', 

35 'GPU_PRECISION': 'float32', 

36 'GPU_BATCHED_MAX': '32', 

37 } 

38 

39 if variable not in defaults.keys(): 

40 allowed = '\n'.join([f' STRONGCOCA_{var}' for var in defaults.keys()]) 

41 raise ValueError(f'Environment variable STRONGCOCA_{variable} is unknown. ' 

42 f'Allowed variables are:\n{allowed}') 

43 

44 return getenv(f'STRONGCOCA_{variable}', defaults[variable]) 

45 

46 

47def get_int(variable) -> int: 

48 strvalue = strongcoca_getenv(variable) 

49 try: 

50 value = int(strvalue) 

51 except ValueError: 

52 raise ValueError(f'Expected environment variable STRONGCOCA_{variable} to ' 

53 f'be castable to int. Value is {strvalue!r}.') 

54 return value 

55 

56 

57def get_float(variable) -> float: 

58 strvalue = strongcoca_getenv(variable) 

59 try: 

60 value = float(strvalue) 

61 except ValueError: 

62 raise ValueError(f'Expected environment variable STRONGCOCA_{variable} to ' 

63 f'be castable to float. Value is {strvalue!r}.') 

64 

65 return value