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

59 statements  

« prev     ^ index     » next       coverage.py v7.4.0, created at 2024-01-03 16:47 -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 "Appends sub to constants, sweep, or linkedsweep." 

37 sweepsub, sweepval = splitsweep(sub) 

38 if sweepsub: # if the whole key is swept 

39 sub = sweepval 

40 for key in keys: 

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

42 value = sub 

43 else: 

44 with pywarnings.catch_warnings(): 

45 pywarnings.filterwarnings("error") 

46 try: 

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

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

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

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

51 if key.shape == sub.shape: 

52 value = sub[key.idx] 

53 sweepel, sweepval = splitsweep(value) 

54 if sweepel: # if only an element is swept 

55 value = sweepval 

56 sweepsub = True 

57 elif sweepsub: 

58 try: 

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

60 except ValueError: 

61 raise ValueError("cannot sweep variable %s of shape %s" 

62 " with array of shape %s; array shape" 

63 " must either be %s or %s" % 

64 (key.veckey, key.shape, sub.shape, 

65 key.shape, ("N",)+key.shape)) 

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

67 value = sub[idx] 

68 else: 

69 raise ValueError("cannot substitute array of shape %s for" 

70 " variable %s of shape %s." % 

71 (sub.shape, key.veckey, key.shape)) 

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

73 linkedsweep[key] = value 

74 elif sweepsub: 

75 sweep[key] = value 

76 else: 

77 try: 

78 assert np.isnan(value) 

79 except (AssertionError, TypeError, ValueError): 

80 constants[key] = value