Coverage for gpaw/wfd_operators.py: 98%

1import gpaw.cgpaw as cgpaw 

2from gpaw.fd_operators import FDOperator 

3import numpy as np 

4 

5 

6class WeightedFDOperator(FDOperator): 

7 def __init__(self, operators): 

8 """Compound operator A with nonlocal weights. 

9 

10 A = Sum_i weights_i * operators_i 

11 

12 Arguments: 

13 operators -- List of FDOperators. 

14 

15 The weights have to be set using the set_weights method. 

16 A is build from operators which are then discarded. 

17 """ 

18 self.gd = operators[0].gd 

19 self.shape = tuple(self.gd.n_c) 

20 self.coef_ps = [] 

21 self.offset_ps = [] 

22 for op in operators: 

23 assert self.gd == op.gd 

24 assert operators[0].dtype == op.dtype 

25 assert not hasattr(op, 'nweights') 

26 if 0 in op.offset_p: 

27 assert op.offset_p[0] == 0 

28 self.offset_ps.append( 

29 np.ascontiguousarray(op.offset_p.copy()) 

30 ) 

31 self.coef_ps.append( 

32 np.ascontiguousarray(op.coef_p.copy()) 

33 ) 

34 else: 

35 self.offset_ps.append( 

36 np.ascontiguousarray(np.hstack(([0], op.offset_p)))) 

37 self.coef_ps.append( 

38 np.ascontiguousarray(np.hstack(([.0], op.coef_p)))) 

39 assert self.offset_ps[-1][0] == 0 

40 assert len(self.coef_ps[-1]) == len(self.offset_ps[-1]) 

41 assert self.offset_ps[-1].flags.c_contiguous 

42 assert self.coef_ps[-1].flags.c_contiguous 

43 assert self.offset_ps[-1].dtype == int 

44 assert self.coef_ps[-1].dtype == float 

45 

46 self.nweights = len(operators) 

47 self.cfd = all([op.cfd for op in operators]) 

48 self.mp = max([op.mp for op in operators]) 

49 self.dtype = operators[0].dtype 

50 if self.gd.comm.size > 1: 

51 self.comm = self.gd.comm.get_c_object() 

52 else: 

53 self.comm = None 

54 self.npoints = max([op.npoints for op in operators]) 

55 self.weights = None 

56 self.operator = None 

57 self.description = 'Weighted Finite Difference Operator\n ' 

58 self.description += '\n '.join([op.description 

59 for op in operators]) 

60 

61 self.xp = np 

62 

63 def set_weights(self, weights): 

64 """Set the operator weights. 

65 

66 weights -- List of numpy arrays sized gd.n_c 

67 (weights are not copied). 

68 """ 

69 assert len(weights) == self.nweights 

70 for weight in weights: 

71 assert weight.shape == self.shape 

72 assert weight.dtype == float 

73 assert weight.flags.c_contiguous 

74 self.weights = weights 

75 self.operator = cgpaw.WOperator( 

76 self.nweights, self.weights, 

77 self.coef_ps, self.offset_ps, self.gd.n_c, self.mp, 

78 self.gd.neighbor_cd, self.dtype == float, self.comm, self.cfd 

79 )