Coverage for gpkit/nomials/substitution.py: 93%

59 statements  

« prev     ^ index     » next       coverage.py v7.4.0, created at 2024-01-07 22:15 -0500

1"Scripts to parse and collate substitutions" 

2import warnings as pywarnings 

3import numpy as np 

4from ..small_scripts import splitsweep 

5from ..keydict import KeySet 

6 

7 

8def parse_subs(varkeys, substitutions, clean=False): 

9 "Seperates subs into the constants, sweeps, linkedsweeps actually present." 

10 constants, sweep, linkedsweep = {}, {}, {} 

11 if clean: 

12 for var in varkeys: 

13 if dict.__contains__(substitutions, var): 

14 sub = dict.__getitem__(substitutions, var) 

15 append_sub(sub, [var], constants, sweep, linkedsweep) 

16 else: 

17 if not hasattr(varkeys, "keymap"): 

18 varkeys = KeySet(varkeys) 

19 varkeys.update_keymap() 

20 if hasattr(substitutions, "keymap"): 

21 for var in varkeys.keymap: 

22 if dict.__contains__(substitutions, var): 

23 sub = dict.__getitem__(substitutions, var) 

24 keys = varkeys.keymap[var] 

25 append_sub(sub, keys, constants, sweep, linkedsweep) 

26 else: 

27 for var in substitutions: 

28 key = getattr(var, "key", var) 

29 if key in varkeys.keymap: 

30 sub, keys = substitutions[var], varkeys.keymap[key] 

31 append_sub(sub, keys, constants, sweep, linkedsweep) 

32 return constants, sweep, linkedsweep 

33 

34 

35def append_sub(sub, keys, constants, sweep, linkedsweep): 

36 # pylint: disable=too-many-branches 

37 "Appends sub to constants, sweep, or linkedsweep." 

38 sweepsub, sweepval = splitsweep(sub) 

39 if sweepsub: # if the whole key is swept 

40 sub = sweepval 

41 for key in keys: 

42 if not key.shape or not getattr(sub, "shape", hasattr(sub, "__len__")): 

43 value = sub 

44 else: 

45 with pywarnings.catch_warnings(): 

46 pywarnings.filterwarnings("error") 

47 try: 

48 sub = np.array(sub) if not hasattr(sub, "shape") else sub 

49 except Warning: # pragma: no cover #TODO: coverage this 

50 # ragged nested sequences, eg [[2]], [3, 4]], in py3.7+ 

51 sub = np.array(sub, dtype=object) 

52 if key.shape == sub.shape: 

53 value = sub[key.idx] 

54 sweepel, sweepval = splitsweep(value) 

55 if sweepel: # if only an element is swept 

56 value = sweepval 

57 sweepsub = True 

58 elif sweepsub: 

59 try: 

60 np.broadcast(sub, np.empty(key.shape)) 

61 except ValueError as exc: 

62 raise ValueError(f"cannot sweep variable {key.veckey} of " 

63 f"shape {key.shape} with array of shape " 

64 f"{sub.shape}; array shape must either be" 

65 f" {key.shape} or {('N',)+key.shape}" 

66 ) from exc 

67 idx = (slice(None),)+key.descr["idx"] 

68 value = sub[idx] 

69 else: 

70 raise ValueError(f"cannot substitute array of shape {sub.shape}" 

71 f" for variable {key.veckey} of shape " 

72 f"{key.shape}.") 

73 if hasattr(value, "__call__") and not hasattr(value, "key"): 

74 linkedsweep[key] = value 

75 elif sweepsub: 

76 sweep[key] = value 

77 else: 

78 try: 

79 assert np.isnan(value) 

80 except (AssertionError, TypeError, ValueError): 

81 constants[key] = value