am_fft_demo.py

#!/usr/bin/env python

from enthought.traits.api import *
from enthought.traits.ui.api import *
from enthought.chaco.api import *

from enthought.enable.component_editor import ComponentEditor
from enthought.chaco.chaco_plot_editor import ChacoPlotItem
from numpy import arange
from numpy import sin, cos, pi
from scipy import fft
from scipy.fftpack import fftshift


class Data(HasTraits):
    tt = Array
    ff = Array
    yy = Property(Array, depends_on=['carrier_amplitude', 'carrier_frequency', 
      'mod_frequency', 'mod_amplitude'])
    amp = Property(Array, depends_on=['carrier_amplitude', 'carrier_frequency', 
      'mod_frequency', 'mod_amplitude'])
    carrier_amplitude         = Range(low=0.0,high=2.0,value=1.0)
    carrier_frequency = Range(low=.01,high=100.0,value=20)
    mod_frequency     = Range(low=.01, high=10.0,value=1)
    mod_amplitude               = Range(low=0.01, high=2, val=0.5)

    plot_type = Enum("line", "scatter")

    traits_view = View(
                       Group(ChacoPlotItem("tt", "yy",
                                                  type_trait='plot_type',
                                                  resizable=True,
                                                  x_label="time",
                                                  y_label="amplitude",
                                                  x_bounds=(0,10),
                                                  x_auto=False,
                                                  y_bounds=(-50,50),
                                                  y_auto=True,
                                                  color="blue",
                                                  bgcolor="white",
                                                  border_visible=True,
                                                  border_width=1,
                                                  title='AM Demo',
                                                  padding_bg_color="lightgray"),
                      ChacoPlotItem("ff", "amp",
                                                  type_trait='plot_type',
                                                  resizable=True,
                                                  x_label="frequency",
                                                  y_label="amplitude",
                                                  x_bounds=(0,10),
                                                  x_auto=True,
                                                  y_bounds=(-50,50),
                                                  y_auto=True,
                                                  color="blue",
                                                  bgcolor="white",
                                                  border_visible=True,
                                                  border_width=1,
                                                  title='AM Demo',
                                                  padding_bg_color="lightgray"),                    
                             Item(name='carrier_amplitude'),
                             Item(name='carrier_frequency'),
                             Item(name='mod_frequency'),
                             Item(name='mod_amplitude'),
                             Item(name='plot_type')
                       ),
                       resizable = True,
                       buttons = ["OK"],
                       title='AM Demo',
                       width=900, height=800)
    
    def _tt_default(self):
        """ Default handler for volume Trait Array. """
        return arange(0, 100, 0.001)

    def _ff_default(self):
        return arange(-50, 50, 0.001)

    def _get_amp(self):
      return (fftshift(abs(fft(self.yy))))
      

    def _get_yy(self):
        """Recalculate when one a trait the property depends on changes."""
        w_c = 2*pi*self.carrier_frequency
        w_m = 2*pi*self.mod_frequency
        
        return ((self.carrier_amplitude + self.mod_amplitude * cos(self.tt*w_m)) *
                sin(self.tt*(w_c)))
                                
                             
      
if __name__ == '__main__':
    viewer = Data()
    viewer.configure_traits()