25 października 2014

generowanie funkcji w matplotlib

Aktualizacja: Nowa wersja klasy: https://blog.kkthx.pl/2015/05/python-dsp-framework/


Klasa ułatwiająca operowanie na dyskretnych funkcjach trygonometrycznych (ustawianie funkcji, amplitudy, fazy, bufora, samplowania; sumowanie, mnożenie funkcji).
Przykład użycia:

trigfn().show()

 

#-*- encoding: utf-8 -*-
# v.1.6
import matplotlib.pylab as mpl
import numpy as np

class trigfn:
	def __init__(self, wave=np.sin, buf=(0, 1), sampl=64, freq=1, ampl=1, phase=0, label=''):
		'''
		* wave - "np.sin", "np.cos" or points list explicite
		* buf - buffer size
			(<from>, <to>)
		* sampl - sampling in Hz
		* freq - frequency in Hz
		* ampl - amplitude
		* phase - phase in degrees
		* label - label for plot/legend

		'''
		self.wave = wave
		self.buf = buf
		self.sampl = sampl
		self.freq = freq
		self.ampl = ampl
		self.phase = phase
		self.label = label

		self.raw_dft = None
		self.real_dft = None
		self.fig = None
		self.pwa = []

		self.inputs = np.arange(self.buf[0], self.buf[1], 1/self.sampl)
		self.fninputs = self.inputs*2*np.pi
		if isinstance(wave, np.ufunc):
			self.fnoutputs = self.ampl*self.wave(self.freq*self.fninputs + (self.phase%360/360)*2*np.pi)
		else:
			self.fnoutputs = wave

	def plot(self, arg=None):
		self.plott(arg)

	def plott(self, arg=None):
		'''
		Plot samples in time domain
		'''
		if (arg is not None):
			arg.grid(True)
			arg.plot(self.fninputs, self.fnoutputs, label=self.label)
		else:
			mpl.grid(True)
			mpl.plot(self.fninputs, self.fnoutputs, label=self.label)

		mpl.plot(self.fninputs, np.zeros(len(self.fninputs)), 'k')
		mpl.margins(0.05)

	def plotf(self):
		'''
		Plot in frequency domain (DFT on samples and then plot)
		'''
		if (self.raw_dft is None):
			self.dft()
		mpl.grid(True)
		mpl.plot(range(0, len(self.real_dft)), self.real_dft, c='gray', ls='-')
		for i, p in enumerate(self.real_dft):
			if (p>1e-10):
				point, = mpl.plot(i, p, 'ro', ms=5)
				if self.fig:
					'''
					http://stackoverflow.com/questions/11537374/matplotlib-basemap-popup-box/11556140#11556140
					we can set self.fig and hover a mouse over a point to see
					annotation with frequency and amplitude in dft graph
					'''
					annotation = mpl.annotate("F=%sHz\nA=%s" % (i, p),
							xy=(i, p), xycoords='data', xytext=(i, p), textcoords='data', horizontalalignment="right",
							bbox=dict(boxstyle="round", facecolor="w", edgecolor="0.5", alpha=0.9))
					annotation.set_visible(False)
					self.pwa.append([point, annotation])

			else:
				mpl.plot(i, p, 'ko', ms=3)
		if self.fig:
			self.fig.canvas.mpl_connect('motion_notify_event', self.on_move)
		mpl.xticks(range(0, len(self.real_dft),  int(self.sampl/32) if int(self.sampl/32)>0 else 1), rotation=0)
		mpl.margins(0.05)
		mpl.xlabel('Frequency [Hz]')
		mpl.ylabel('Amplitude')

	def show(self, title=None):
		if (title is not None):
			mpl.title(title)
		self.plot()
		#float unstability errors fix
		#mpl.ylim(1e-6)
		mpl.show()

	def dft(self):
		self.raw_dft = np.zeros(np.size(self.fnoutputs), dtype=complex)
		N = np.size(self.fninputs)
		for k in range(N):
			for n, x in enumerate(self.fnoutputs):
				self.raw_dft[k] += x * (np.cos((2*np.pi*k*n)/N)-1j*np.sin((2*np.pi*k*n)/N))

		self.real_dft = np.absolute(self.raw_dft[0:len(self.raw_dft)/2])/N*2

	def __add__(self, other):
		if isinstance(other, trigfn):
			return trigfn(wave=self.fnoutputs+other.fnoutputs, buf=self.buf, sampl=self.sampl)
		else:
			return trigfn(wave=self.fnoutputs+other, buf=self.buf, sampl=self.sampl)

	def __sub__(self, other):
		if isinstance(other, trigfn):
			return trigfn(wave=self.fnoutputs-other.fnoutputs, buf=self.buf, sampl=self.sampl)
		else:
			return trigfn(wave=self.fnoutputs-other, buf=self.buf, sampl=self.sampl)

	def __mul__(self, other):
		if isinstance(other, trigfn):
			return trigfn(wave=self.fnoutputs*other.fnoutputs, buf=self.buf, sampl=self.sampl)
		else:
			return trigfn(wave=self.fnoutputs*other, buf=self.buf, sampl=self.sampl)

	def __truediv__(self, other):
		if isinstance(other, trigfn):
			return trigfn(wave=self.fnoutputs/other.fnoutputs, buf=self.buf, sampl=self.sampl)
		else:
			return trigfn(wave=self.fnoutputs/other, buf=self.buf, sampl=self.sampl)

	def on_move(self, event):
		for point, annotation in self.pwa:
			should_be_visible = (point.contains(event)[0] == True)
			if should_be_visible != annotation.get_visible():
				annotation.set_visible(should_be_visible)
				mpl.draw()

TrackBack

TrackBack URL dla tej wiadomości:
https://blog.kkthx.pl/2014/10/generowanie-funkcji-w-matplotlib/trackback/

Napisz komentarz