def subplot_xy(ax, x, y, **kwargs): if is_np_array_type(y.__class__): if y.ndim > 1: for i in range(y.shape[0]): subplot_xy(ax, x, y[i], **kwargs) else: ax.plot(x, y, **kwargs) elif isinstance(y, list): for i in range(y.__len__()): subplot_xy(ax, x, y[i], **kwargs) else: ax.plot(x, y, **kwargs)
def W(p=None, w=None, k=None, fc=100e3, c=1480, N=500): if p==None: error('\'p\' are missing. Aborting.') return if k!=None: k_abs = k elif p!=None and fc!=None and c!=None: k_abs = 2*pi*fc/c else: error('Either \'k\', or \'c\' and \'fc\' must be supplied. Aborting.') return theta = np.linspace(-0.5,0.5,N)*pi k = Coordinate(r=k_abs, theta=theta, phi=None) def compute_W(ax,ay,az, w): Wx = dot( ax, w ) Wy = dot( ay, w ) Wz = dot( az, w ) Wxyz = vstack((Wx,Wy,Wz)).T return Wxyz # pT = p.T.copy() ax = exp(1j*outer( k[:,0], p[:,0])) ay = exp(1j*outer( k[:,1], p[:,1])) az = exp(1j*outer( k[:,2], p[:,2])) if type(w) == list: W = [] for wi in w: W.append(compute_W(ax,ay,az, wi)) elif is_np_array_type(w.__class__): if w.ndim == 1: W = [] W.append( compute_W(ax,ay,az,w) ) elif w.ndim == 2: W1 = compute_W(ax,ay,az,w[0]) W = np.zeros((w.shape[0],W1.shape[0],W1.shape[1]),dtype=W1.dtype) W[0] = W1 for i in range(1,w.shape[0]): W[i] = compute_W(ax,ay,az, w[i]) elif w.ndim == 3: W1 = compute_W(ax,ay,az,w[0,0]) # Not exactly efficient, but shouldn't be noticeable W = np.zeros((w.shape[0],w.shape[1],W1.shape[0],W1.shape[1]),dtype=W1.dtype) for i in range(0,w.shape[0]): for j in range(0,w.shape[1]): W[i,j] = compute_W(ax,ay,az, w[i,j]) else: print 'EE: Not implemented.' else: W = [] W.append( compute_W(ax,ay,az,w) ) return W