<s> """<STR_LIT>""" <EOL> import numpy as np <EOL> from traits . api import Bool , HasTraits , Instance , on_trait_change <EOL> from tvtk . api import tvtk <EOL> from tvtk . common import camel2enthought <EOL> from mayavi . sources . array_source import ArraySource <EOL> from mayavi . core . registry import registry <EOL> from mayavi . core . trait_defs import ArrayNumberOrNone , ArrayOrNone <EOL> from . import tools <EOL> from . engine_manager import get_null_engine , engine_manager <EOL> __all__ = [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , <EOL> '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , <EOL> '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , <EOL> ] <EOL> class MlabSource ( HasTraits ) : <EOL> """<STR_LIT>""" <EOL> dataset = Instance ( tvtk . DataSet ) <EOL> m_data = Instance ( HasTraits ) <EOL> _disable_update = Bool ( False ) <EOL> def reset ( self , ** traits ) : <EOL> """<STR_LIT>""" <EOL> raise NotImplementedError ( ) <EOL> def update ( self ) : <EOL> """<STR_LIT>""" <EOL> if not self . _disable_update : <EOL> self . dataset . modified ( ) <EOL> md = self . m_data <EOL> if md is not None : <EOL> if hasattr ( md , '<STR_LIT>' ) : <EOL> md . _assign_attribute . update ( ) <EOL> md . data_changed = True <EOL> def set ( self , trait_change_notify = True , ** traits ) : <EOL> """<STR_LIT>""" <EOL> try : <EOL> self . _disable_update = True <EOL> super ( MlabSource , self ) . set ( trait_change_notify , ** traits ) <EOL> finally : <EOL> self . _disable_update = False <EOL> if trait_change_notify : <EOL> self . update ( ) <EOL> return self <EOL> def _m_data_changed ( self , ds ) : <EOL> if not hasattr ( ds , '<STR_LIT>' ) : <EOL> ds . add_trait ( '<STR_LIT>' , Instance ( MlabSource ) ) <EOL> ds . mlab_source = self <EOL> class MGlyphSource ( MlabSource ) : <EOL> """<STR_LIT>""" <EOL> x = ArrayNumberOrNone <EOL> y = ArrayNumberOrNone <EOL> z = ArrayNumberOrNone <EOL> points = ArrayOrNone <EOL> scalars = ArrayNumberOrNone <EOL> u = ArrayNumberOrNone <EOL> v = ArrayNumberOrNone <EOL> w = ArrayNumberOrNone <EOL> vectors = ArrayOrNone <EOL> def reset ( self , ** traits ) : <EOL> """<STR_LIT>""" <EOL> for name in ( '<STR_LIT:x>' , '<STR_LIT:y>' , '<STR_LIT:z>' , '<STR_LIT:u>' , '<STR_LIT:v>' , '<STR_LIT:w>' , '<STR_LIT>' ) : <EOL> if name in traits and traits [ name ] is not None : <EOL> traits [ name ] = np . atleast_1d ( traits [ name ] ) <EOL> self . set ( trait_change_notify = False , ** traits ) <EOL> vectors = self . vectors <EOL> scalars = self . scalars <EOL> points = self . points <EOL> x , y , z = self . x , self . y , self . z <EOL> x = np . atleast_1d ( x ) <EOL> y = np . atleast_1d ( y ) <EOL> z = np . atleast_1d ( z ) <EOL> if '<STR_LIT>' in traits : <EOL> x = points [ : , <NUM_LIT:0> ] . ravel ( ) <EOL> y = points [ : , <NUM_LIT:1> ] . ravel ( ) <EOL> z = points [ : , <NUM_LIT:2> ] . ravel ( ) <EOL> self . set ( x = x , y = y , z = z , trait_change_notify = False ) <EOL> else : <EOL> points = np . c_ [ x . ravel ( ) , y . ravel ( ) , z . ravel ( ) ] . ravel ( ) <EOL> points . shape = ( - <NUM_LIT:1> , <NUM_LIT:3> ) <EOL> self . set ( points = points , trait_change_notify = False ) <EOL> u , v , w = self . u , self . v , self . w <EOL> if u is not None : <EOL> u = np . atleast_1d ( u ) <EOL> v = np . atleast_1d ( v ) <EOL> w = np . atleast_1d ( w ) <EOL> if len ( u ) > <NUM_LIT:0> : <EOL> vectors = np . c_ [ u . ravel ( ) , v . ravel ( ) , <EOL> w . ravel ( ) ] . ravel ( ) <EOL> vectors . shape = ( - <NUM_LIT:1> , <NUM_LIT:3> ) <EOL> self . set ( vectors = vectors , trait_change_notify = False ) <EOL> if '<STR_LIT>' in traits : <EOL> u = vectors [ : , <NUM_LIT:0> ] . ravel ( ) <EOL> v = vectors [ : , <NUM_LIT:1> ] . ravel ( ) <EOL> w = vectors [ : , <NUM_LIT:2> ] . ravel ( ) <EOL> self . set ( u = u , v = v , w = w , trait_change_notify = False ) <EOL> else : <EOL> if u is not None and len ( u ) > <NUM_LIT:0> : <EOL> vectors = np . c_ [ u . ravel ( ) , v . ravel ( ) , <EOL> w . ravel ( ) ] . ravel ( ) <EOL> vectors . shape = ( - <NUM_LIT:1> , <NUM_LIT:3> ) <EOL> self . set ( vectors = vectors , trait_change_notify = False ) <EOL> if vectors is not None and len ( vectors ) > <NUM_LIT:0> : <EOL> assert len ( points ) == len ( vectors ) <EOL> if scalars is not None : <EOL> scalars = np . atleast_1d ( scalars ) <EOL> if len ( scalars ) > <NUM_LIT:0> : <EOL> assert len ( points ) == len ( scalars . ravel ( ) ) <EOL> polys = np . arange ( <NUM_LIT:0> , len ( points ) , <NUM_LIT:1> , '<STR_LIT:l>' ) <EOL> polys = np . reshape ( polys , ( len ( points ) , <NUM_LIT:1> ) ) <EOL> if self . dataset is None : <EOL> pd = tvtk . PolyData ( ) <EOL> else : <EOL> pd = self . dataset <EOL> pd . set ( points = points , polys = polys ) <EOL> if self . vectors is not None : <EOL> pd . point_data . vectors = self . vectors <EOL> pd . point_data . vectors . name = '<STR_LIT>' <EOL> if self . scalars is not None : <EOL> pd . point_data . scalars = self . scalars . ravel ( ) <EOL> pd . point_data . scalars . name = '<STR_LIT>' <EOL> self . dataset = pd <EOL> def _x_changed ( self , x ) : <EOL> x = np . atleast_1d ( x ) <EOL> self . points [ : , <NUM_LIT:0> ] = x . ravel ( ) <EOL> self . update ( ) <EOL> def _y_changed ( self , y ) : <EOL> y = np . atleast_1d ( y ) <EOL> self . points [ : , <NUM_LIT:1> ] = y . ravel ( ) <EOL> self . update ( ) <EOL> def _z_changed ( self , z ) : <EOL> z = np . atleast_1d ( z ) <EOL> self . points [ : , <NUM_LIT:2> ] = z . ravel ( ) <EOL> self . update ( ) <EOL> def _u_changed ( self , u ) : <EOL> u = np . atleast_1d ( u ) <EOL> self . vectors [ : , <NUM_LIT:0> ] = u . ravel ( ) <EOL> self . update ( ) <EOL> def _v_changed ( self , v ) : <EOL> v = np . atleast_1d ( v ) <EOL> self . vectors [ : , <NUM_LIT:1> ] = v . ravel ( ) <EOL> self . update ( ) <EOL> def _w_changed ( self , w ) : <EOL> w = np . atleast_1d ( w ) <EOL> self . vectors [ : , <NUM_LIT:2> ] = w . ravel ( ) <EOL> self . update ( ) <EOL> def _points_changed ( self , p ) : <EOL> p = np . atleast_2d ( p ) <EOL> self . dataset . points = p <EOL> self . update ( ) <EOL> def _scalars_changed ( self , s ) : <EOL> if s is None : <EOL> self . dataset . point_data . scalars = None <EOL> self . dataset . point_data . remove_array ( '<STR_LIT>' ) <EOL> else : <EOL> s = np . atleast_1d ( s ) <EOL> self . dataset . point_data . scalars = s . ravel ( ) <EOL> self . dataset . point_data . scalars . name = '<STR_LIT>' <EOL> self . update ( ) <EOL> def _vectors_changed ( self , v ) : <EOL> self . dataset . point_data . vectors = v <EOL> self . dataset . point_data . vectors . name = '<STR_LIT>' <EOL> self . update ( ) <EOL> class MVerticalGlyphSource ( MGlyphSource ) : <EOL> """<STR_LIT>""" <EOL> def reset ( self , ** traits ) : <EOL> """<STR_LIT>""" <EOL> if '<STR_LIT>' in traits : <EOL> s = traits [ '<STR_LIT>' ] <EOL> if s is not None : <EOL> traits [ '<STR_LIT:u>' ] = traits [ '<STR_LIT:v>' ] = np . ones_like ( s ) , <EOL> traits [ '<STR_LIT:w>' ] = s <EOL> super ( MVerticalGlyphSource , self ) . reset ( ** traits ) <EOL> def _scalars_changed ( self , s ) : <EOL> self . dataset . point_data . scalars = s <EOL> self . dataset . point_data . scalars . name = '<STR_LIT>' <EOL> self . set ( vectors = np . c_ [ np . ones_like ( s ) , <EOL> np . ones_like ( s ) , <EOL> s ] ) <EOL> self . update ( ) <EOL> class MArraySource ( MlabSource ) : <EOL> """<STR_LIT>""" <EOL> x = ArrayOrNone <EOL> y = ArrayOrNone <EOL> z = ArrayOrNone <EOL> scalars = ArrayOrNone <EOL> u = ArrayOrNone <EOL> v = ArrayOrNone <EOL> w = ArrayOrNone <EOL> vectors = ArrayOrNone <EOL> def reset ( self , ** traits ) : <EOL> """<STR_LIT>""" <EOL> self . set ( trait_change_notify = False , ** traits ) <EOL> vectors = self . vectors <EOL> scalars = self . scalars <EOL> x , y , z = [ np . atleast_3d ( a ) for a in ( self . x , self . y , self . z ) ] <EOL> u , v , w = self . u , self . v , self . w <EOL> if '<STR_LIT>' in traits : <EOL> u = vectors [ : , <NUM_LIT:0> ] . ravel ( ) <EOL> v = vectors [ : , <NUM_LIT:1> ] . ravel ( ) <EOL> w = vectors [ : , <NUM_LIT:2> ] . ravel ( ) <EOL> self . set ( u = u , v = v , w = w , trait_change_notify = False ) <EOL> else : <EOL> if u is not None and len ( u ) > <NUM_LIT:0> : <EOL> vectors = np . c_ [ u . ravel ( ) , v . ravel ( ) , <EOL> w . ravel ( ) ] . ravel ( ) <EOL> vectors . shape = ( u . shape [ <NUM_LIT:0> ] , u . shape [ <NUM_LIT:1> ] , w . shape [ <NUM_LIT:2> ] , <NUM_LIT:3> ) <EOL> self . set ( vectors = vectors , trait_change_notify = False ) <EOL> if vectors is not None and len ( vectors ) > <NUM_LIT:0> and scalars is not None : <EOL> assert len ( scalars ) == len ( vectors ) <EOL> if x . shape [ <NUM_LIT:0> ] <= <NUM_LIT:1> : <EOL> dx = <NUM_LIT:1> <EOL> else : <EOL> dx = x [ <NUM_LIT:1> , <NUM_LIT:0> , <NUM_LIT:0> ] - x [ <NUM_LIT:0> , <NUM_LIT:0> , <NUM_LIT:0> ] <EOL> if y . shape [ <NUM_LIT:1> ] <= <NUM_LIT:1> : <EOL> dy = <NUM_LIT:1> <EOL> else : <EOL> dy = y [ <NUM_LIT:0> , <NUM_LIT:1> , <NUM_LIT:0> ] - y [ <NUM_LIT:0> , <NUM_LIT:0> , <NUM_LIT:0> ] <EOL> if z . shape [ <NUM_LIT:2> ] <= <NUM_LIT:1> : <EOL> dz = <NUM_LIT:1> <EOL> else : <EOL> dz = z [ <NUM_LIT:0> , <NUM_LIT:0> , <NUM_LIT:1> ] - z [ <NUM_LIT:0> , <NUM_LIT:0> , <NUM_LIT:0> ] <EOL> if self . m_data is None : <EOL> ds = ArraySource ( transpose_input_array = True ) <EOL> else : <EOL> ds = self . m_data <EOL> old_scalar = ds . scalar_data <EOL> ds . set ( vector_data = vectors , <EOL> origin = [ x . min ( ) , y . min ( ) , z . min ( ) ] , <EOL> spacing = [ dx , dy , dz ] , <EOL> scalar_data = scalars ) <EOL> if scalars is old_scalar : <EOL> ds . _scalar_data_changed ( scalars ) <EOL> self . dataset = ds . image_data <EOL> self . m_data = ds <EOL> @ on_trait_change ( '<STR_LIT>' ) <EOL> def _xyz_changed ( self ) : <EOL> x , y , z = self . x , self . y , self . z <EOL> dx = x [ <NUM_LIT:1> , <NUM_LIT:0> , <NUM_LIT:0> ] - x [ <NUM_LIT:0> , <NUM_LIT:0> , <NUM_LIT:0> ] <EOL> dy = y [ <NUM_LIT:0> , <NUM_LIT:1> , <NUM_LIT:0> ] - y [ <NUM_LIT:0> , <NUM_LIT:0> , <NUM_LIT:0> ] <EOL> dz = z [ <NUM_LIT:0> , <NUM_LIT:0> , <NUM_LIT:1> ] - z [ <NUM_LIT:0> , <NUM_LIT:0> , <NUM_LIT:0> ] <EOL> ds = self . dataset <EOL> ds . origin = [ x . min ( ) , y . min ( ) , z . min ( ) ] <EOL> ds . spacing = [ dx , dy , dz ] <EOL> if self . m_data is not None : <EOL> self . m_data . set ( origin = ds . origin , spacing = ds . spacing ) <EOL> self . update ( ) <EOL> def _u_changed ( self , u ) : <EOL> self . vectors [ ... , <NUM_LIT:0> ] = u <EOL> self . m_data . _vector_data_changed ( self . vectors ) <EOL> def _v_changed ( self , v ) : <EOL> self . vectors [ ... , <NUM_LIT:1> ] = v <EOL> self . m_data . _vector_data_changed ( self . vectors ) <EOL> def _w_changed ( self , w ) : <EOL> self . vectors [ ... , <NUM_LIT:2> ] = w <EOL> self . m_data . _vector_data_changed ( self . vectors ) <EOL> def _scalars_changed ( self , s ) : <EOL> old = self . m_data . scalar_data <EOL> self . m_data . scalar_data = s <EOL> if old is s : <EOL> self . m_data . _scalar_data_changed ( s ) <EOL> def _vectors_changed ( self , v ) : <EOL> self . m_data . vector_data = v <EOL> class MLineSource ( MlabSource ) : <EOL> """<STR_LIT>""" <EOL> x = ArrayOrNone <EOL> y = ArrayOrNone <EOL> z = ArrayOrNone <EOL> points = ArrayOrNone <EOL> scalars = ArrayOrNone <EOL> def reset ( self , ** traits ) : <EOL> """<STR_LIT>""" <EOL> self . set ( trait_change_notify = False , ** traits ) <EOL> points = self . points <EOL> scalars = self . scalars <EOL> x , y , z = self . x , self . y , self . z <EOL> if '<STR_LIT>' in traits : <EOL> x = points [ : , <NUM_LIT:0> ] . ravel ( ) <EOL> y = points [ : , <NUM_LIT:1> ] . ravel ( ) <EOL> z = points [ : , <NUM_LIT:2> ] . ravel ( ) <EOL> self . set ( x = x , y = y , z = z , trait_change_notify = False ) <EOL> else : <EOL> points = np . c_ [ x . ravel ( ) , y . ravel ( ) , z . ravel ( ) ] . ravel ( ) <EOL> points . shape = ( len ( x ) , <NUM_LIT:3> ) <EOL> self . set ( points = points , trait_change_notify = False ) <EOL> n_pts = len ( points ) - <NUM_LIT:1> <EOL> lines = np . zeros ( ( n_pts , <NUM_LIT:2> ) , '<STR_LIT:l>' ) <EOL> lines [ : , <NUM_LIT:0> ] = np . arange ( <NUM_LIT:0> , n_pts - <NUM_LIT:0.5> , <NUM_LIT:1> , '<STR_LIT:l>' ) <EOL> lines [ : , <NUM_LIT:1> ] = np . arange ( <NUM_LIT:1> , n_pts + <NUM_LIT:0.5> , <NUM_LIT:1> , '<STR_LIT:l>' ) <EOL> if self . dataset is None : <EOL> pd = tvtk . PolyData ( ) <EOL> else : <EOL> pd = self . dataset <EOL> pd . set ( lines = None ) <EOL> pd . set ( points = points ) <EOL> pd . set ( lines = lines ) <EOL> if scalars is not None and len ( scalars ) > <NUM_LIT:0> : <EOL> assert len ( x ) == len ( scalars ) <EOL> pd . point_data . scalars = np . ravel ( scalars ) <EOL> pd . point_data . scalars . name = '<STR_LIT>' <EOL> self . dataset = pd <EOL> def _x_changed ( self , x ) : <EOL> self . points [ : , <NUM_LIT:0> ] = x <EOL> self . update ( ) <EOL> def _y_changed ( self , y ) : <EOL> self . points [ : , <NUM_LIT:1> ] = y <EOL> self . update ( ) <EOL> def _z_changed ( self , z ) : <EOL> self . points [ : , <NUM_LIT:2> ] = z <EOL> self . update ( ) <EOL> def _points_changed ( self , p ) : <EOL> self . dataset . points = p <EOL> self . update ( ) <EOL> def _scalars_changed ( self , s ) : <EOL> self . dataset . point_data . scalars = s . ravel ( ) <EOL> self . dataset . point_data . scalars . name = '<STR_LIT>' <EOL> self . update ( ) <EOL> class MArray2DSource ( MlabSource ) : <EOL> """<STR_LIT>""" <EOL> x = ArrayOrNone <EOL> y = ArrayOrNone <EOL> scalars = ArrayOrNone <EOL> mask = ArrayOrNone <EOL> def reset ( self , ** traits ) : <EOL> """<STR_LIT>""" <EOL> self . set ( trait_change_notify = False , ** traits ) <EOL> x , y , mask = self . x , self . y , self . mask <EOL> scalars = self . scalars <EOL> nx , ny = scalars . shape <EOL> if x is None and y is None : <EOL> x , y = np . mgrid [ - nx / <NUM_LIT> : nx / <NUM_LIT:2> , - ny / <NUM_LIT> : ny / <NUM_LIT:2> ] <EOL> if mask is not None and len ( mask ) > <NUM_LIT:0> : <EOL> scalars [ mask . astype ( '<STR_LIT:bool>' ) ] = np . nan <EOL> scalars = scalars . astype ( '<STR_LIT:float>' ) <EOL> self . set ( scalars = scalars , trait_change_notify = False ) <EOL> z = np . array ( [ <NUM_LIT:0> ] ) <EOL> self . set ( x = x , y = y , z = z , trait_change_notify = False ) <EOL> x = np . atleast_2d ( x . squeeze ( ) . T ) [ <NUM_LIT:0> , : ] . squeeze ( ) <EOL> y = np . atleast_2d ( y . squeeze ( ) ) [ <NUM_LIT:0> , : ] . squeeze ( ) <EOL> if x . ndim == <NUM_LIT:0> : <EOL> dx = <NUM_LIT:1> <EOL> else : <EOL> dx = x [ <NUM_LIT:1> ] - x [ <NUM_LIT:0> ] <EOL> if y . ndim == <NUM_LIT:0> : <EOL> dy = <NUM_LIT:1> <EOL> else : <EOL> dy = y [ <NUM_LIT:1> ] - y [ <NUM_LIT:0> ] <EOL> if self . m_data is None : <EOL> ds = ArraySource ( transpose_input_array = True ) <EOL> else : <EOL> ds = self . m_data <EOL> old_scalar = ds . scalar_data <EOL> ds . set ( origin = [ x . min ( ) , y . min ( ) , <NUM_LIT:0> ] , <EOL> spacing = [ dx , dy , <NUM_LIT:1> ] , <EOL> scalar_data = scalars ) <EOL> if old_scalar is scalars : <EOL> ds . _scalar_data_changed ( scalars ) <EOL> self . dataset = ds . outputs [ <NUM_LIT:0> ] <EOL> self . m_data = ds <EOL> @ on_trait_change ( '<STR_LIT>' ) <EOL> def _xy_changed ( self ) : <EOL> x , y , scalars = self . x , self . y , self . scalars <EOL> nx , ny = scalars . shape <EOL> if x is None or y is None : <EOL> x , y = np . mgrid [ - nx / <NUM_LIT> : nx / <NUM_LIT:2> , - ny / <NUM_LIT> : ny / <NUM_LIT:2> ] <EOL> self . trait_setq ( x = x , y = y ) <EOL> x = np . atleast_2d ( x . squeeze ( ) . T ) [ <NUM_LIT:0> , : ] . squeeze ( ) <EOL> y = np . atleast_2d ( y . squeeze ( ) ) [ <NUM_LIT:0> , : ] . squeeze ( ) <EOL> dx = x [ <NUM_LIT:1> ] - x [ <NUM_LIT:0> ] <EOL> dy = y [ <NUM_LIT:1> ] - y [ <NUM_LIT:0> ] <EOL> ds = self . dataset <EOL> ds . origin = [ x . min ( ) , y . min ( ) , <NUM_LIT:0> ] <EOL> ds . spacing = [ dx , dy , <NUM_LIT:1> ] <EOL> if self . m_data is not None : <EOL> self . m_data . set ( origin = ds . origin , spacing = ds . spacing ) <EOL> self . update ( ) <EOL> def _scalars_changed ( self , s ) : <EOL> mask = self . mask <EOL> if mask is not None and len ( mask ) > <NUM_LIT:0> : <EOL> s [ mask . astype ( '<STR_LIT:bool>' ) ] = np . nan <EOL> s = s . astype ( '<STR_LIT:float>' ) <EOL> self . set ( scalars = s , trait_change_notify = False ) <EOL> old = self . m_data . scalar_data <EOL> self . m_data . scalar_data = s <EOL> if s is old : <EOL> self . m_data . _scalar_data_changed ( s ) <EOL> class MGridSource ( MlabSource ) : <EOL> """<STR_LIT>""" <EOL> x = ArrayOrNone <EOL> y = ArrayOrNone <EOL> z = ArrayOrNone <EOL> points = ArrayOrNone <EOL> scalars = ArrayOrNone <EOL> mask = ArrayOrNone <EOL> def reset ( self , ** traits ) : <EOL> """<STR_LIT>""" <EOL> self . set ( trait_change_notify = False , ** traits ) <EOL> points = self . points <EOL> scalars = self . scalars <EOL> x , y , z , mask = self . x , self . y , self . z , self . mask <EOL> if mask is not None and len ( mask ) > <NUM_LIT:0> : <EOL> scalars [ mask . astype ( '<STR_LIT:bool>' ) ] = np . nan <EOL> scalars = scalars . astype ( '<STR_LIT:float>' ) <EOL> self . set ( scalars = scalars , trait_change_notify = False ) <EOL> assert len ( x . shape ) == <NUM_LIT:2> , "<STR_LIT>" <EOL> assert len ( y . shape ) == <NUM_LIT:2> , "<STR_LIT>" <EOL> assert len ( z . shape ) == <NUM_LIT:2> , "<STR_LIT>" <EOL> assert x . shape == y . shape , "<STR_LIT>" <EOL> assert y . shape == z . shape , "<STR_LIT>" <EOL> nx , ny = x . shape <EOL> points = np . c_ [ x . ravel ( ) , y . ravel ( ) , z . ravel ( ) ] . ravel ( ) <EOL> points . shape = ( nx * ny , <NUM_LIT:3> ) <EOL> self . set ( points = points , trait_change_notify = False ) <EOL> i , j = np . mgrid [ <NUM_LIT:0> : nx - <NUM_LIT:1> , <NUM_LIT:0> : ny - <NUM_LIT:1> ] <EOL> i , j = np . ravel ( i ) , np . ravel ( j ) <EOL> t1 = i * ny + j , ( i + <NUM_LIT:1> ) * ny + j , ( i + <NUM_LIT:1> ) * ny + ( j + <NUM_LIT:1> ) <EOL> t2 = ( i + <NUM_LIT:1> ) * ny + ( j + <NUM_LIT:1> ) , i * ny + ( j + <NUM_LIT:1> ) , i * ny + j <EOL> nt = len ( t1 [ <NUM_LIT:0> ] ) <EOL> triangles = np . zeros ( ( nt * <NUM_LIT:2> , <NUM_LIT:3> ) , '<STR_LIT:l>' ) <EOL> triangles [ <NUM_LIT:0> : nt , <NUM_LIT:0> ] , triangles [ <NUM_LIT:0> : nt , <NUM_LIT:1> ] , triangles [ <NUM_LIT:0> : nt , <NUM_LIT:2> ] = t1 <EOL> triangles [ nt : , <NUM_LIT:0> ] , triangles [ nt : , <NUM_LIT:1> ] , triangles [ nt : , <NUM_LIT:2> ] = t2 <EOL> if self . dataset is None : <EOL> pd = tvtk . PolyData ( ) <EOL> else : <EOL> pd = self . dataset <EOL> pd . set ( points = points , polys = triangles ) <EOL> if scalars is not None and len ( scalars ) > <NUM_LIT:0> : <EOL> if not scalars . flags . contiguous : <EOL> scalars = scalars . copy ( ) <EOL> self . set ( scalars = scalars , trait_change_notify = False ) <EOL> assert x . shape == scalars . shape <EOL> pd . point_data . scalars = scalars . ravel ( ) <EOL> pd . point_data . scalars . name = '<STR_LIT>' <EOL> self . dataset = pd <EOL> def _x_changed ( self , x ) : <EOL> self . trait_setq ( x = x ) <EOL> self . points [ : , <NUM_LIT:0> ] = x . ravel ( ) <EOL> self . update ( ) <EOL> def _y_changed ( self , y ) : <EOL> self . trait_setq ( y = y ) <EOL> self . points [ : , <NUM_LIT:1> ] = y . ravel ( ) <EOL> self . update ( ) <EOL> def _z_changed ( self , z ) : <EOL> self . trait_setq ( z = z ) <EOL> self . points [ : , <NUM_LIT:2> ] = z . ravel ( ) <EOL> self . update ( ) <EOL> def _points_changed ( self , p ) : <EOL> self . dataset . points = p <EOL> self . update ( ) <EOL> def _scalars_changed ( self , s ) : <EOL> mask = self . mask <EOL> if mask is not None and len ( mask ) > <NUM_LIT:0> : <EOL> s [ mask . astype ( '<STR_LIT:bool>' ) ] = np . nan <EOL> s = s . astype ( '<STR_LIT:float>' ) <EOL> self . set ( scalars = s , trait_change_notify = False ) <EOL> self . dataset . point_data . scalars = s . ravel ( ) <EOL> self . dataset . point_data . scalars . name = '<STR_LIT>' <EOL> self . update ( ) <EOL> class MTriangularMeshSource ( MlabSource ) : <EOL> """<STR_LIT>""" <EOL> x = ArrayOrNone <EOL> y = ArrayOrNone <EOL> z = ArrayOrNone <EOL> points = ArrayOrNone <EOL> triangles = ArrayOrNone <EOL> scalars = ArrayOrNone <EOL> def reset ( self , ** traits ) : <EOL> """<STR_LIT>""" <EOL> self . set ( trait_change_notify = False , ** traits ) <EOL> points = self . points <EOL> scalars = self . scalars <EOL> x , y , z = self . x , self . y , self . z <EOL> points = np . c_ [ x . ravel ( ) , y . ravel ( ) , z . ravel ( ) ] . ravel ( ) <EOL> points . shape = ( - <NUM_LIT:1> , <NUM_LIT:3> ) <EOL> self . set ( points = points , trait_change_notify = False ) <EOL> triangles = self . triangles <EOL> assert triangles . shape [ <NUM_LIT:1> ] == <NUM_LIT:3> , "<STR_LIT>" <EOL> assert triangles . max ( ) < len ( points ) , "<STR_LIT>" <EOL> assert triangles . min ( ) >= <NUM_LIT:0> , "<STR_LIT>" <EOL> if self . dataset is None : <EOL> pd = tvtk . PolyData ( ) <EOL> else : <EOL> pd = self . dataset <EOL> pd . set ( points = points ) <EOL> pd . set ( polys = triangles ) <EOL> if ( not '<STR_LIT>' in traits <EOL> and scalars is not None <EOL> and scalars . shape != x . shape ) : <EOL> scalars = z <EOL> if scalars is not None and len ( scalars ) > <NUM_LIT:0> : <EOL> if not scalars . flags . contiguous : <EOL> scalars = scalars . copy ( ) <EOL> self . set ( scalars = scalars , trait_change_notify = False ) <EOL> assert x . shape == scalars . shape <EOL> pd . point_data . scalars = scalars . ravel ( ) <EOL> pd . point_data . scalars . name = '<STR_LIT>' <EOL> self . dataset = pd <EOL> def _x_changed ( self , x ) : <EOL> self . trait_setq ( x = x ) <EOL> self . points [ : , <NUM_LIT:0> ] = x . ravel ( ) <EOL> self . update ( ) <EOL> def _y_changed ( self , y ) : <EOL> self . trait_setq ( y = y ) <EOL> self . points [ : , <NUM_LIT:1> ] = y . ravel ( ) <EOL> self . update ( ) <EOL> def _z_changed ( self , z ) : <EOL> self . trait_setq ( z = z ) <EOL> self . points [ : , <NUM_LIT:2> ] = z . ravel ( ) <EOL> self . update ( ) <EOL> def _points_changed ( self , p ) : <EOL> self . dataset . points = p <EOL> self . update ( ) <EOL> def _scalars_changed ( self , s ) : <EOL> self . dataset . point_data . scalars = s . ravel ( ) <EOL> self . dataset . point_data . scalars . name = '<STR_LIT>' <EOL> self . update ( ) <EOL> def _triangles_changed ( self , triangles ) : <EOL> if triangles . min ( ) < <NUM_LIT:0> : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> if triangles . max ( ) > self . x . size : <EOL> raise ValueError ( '<STR_LIT>' '<STR_LIT>' ) <EOL> self . dataset . polys = triangles <EOL> self . update ( ) <EOL> def convert_to_arrays ( args ) : <EOL> """<STR_LIT>""" <EOL> args = list ( args ) <EOL> for index , arg in enumerate ( args ) : <EOL> if not callable ( arg ) : <EOL> if not hasattr ( arg , '<STR_LIT>' ) : <EOL> arg = np . array ( arg ) <EOL> arg = np . atleast_1d ( arg ) <EOL> if np . any ( np . isinf ( arg ) ) : <EOL> raise ValueError ( """<STR_LIT>""" ) <EOL> args [ index ] = arg <EOL> return args <EOL> def process_regular_vectors ( * args ) : <EOL> """<STR_LIT>""" <EOL> args = convert_to_arrays ( args ) <EOL> if len ( args ) == <NUM_LIT:3> : <EOL> u , v , w = [ np . atleast_3d ( a ) for a in args ] <EOL> assert len ( u . shape ) == <NUM_LIT:3> , "<STR_LIT>" <EOL> x , y , z = np . indices ( u . shape ) <EOL> elif len ( args ) == <NUM_LIT:6> : <EOL> x , y , z , u , v , w = args <EOL> elif len ( args ) == <NUM_LIT:4> : <EOL> x , y , z , f = args <EOL> if not callable ( f ) : <EOL> raise ValueError ( "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> u , v , w = f ( x , y , z ) <EOL> else : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> assert ( x . shape == y . shape and <EOL> y . shape == z . shape and <EOL> u . shape == z . shape and <EOL> v . shape == u . shape and <EOL> w . shape == v . shape ) , "<STR_LIT>" <EOL> return x , y , z , u , v , w <EOL> def process_regular_scalars ( * args ) : <EOL> """<STR_LIT>""" <EOL> args = convert_to_arrays ( args ) <EOL> if len ( args ) == <NUM_LIT:1> : <EOL> s = np . atleast_3d ( args [ <NUM_LIT:0> ] ) <EOL> assert len ( s . shape ) == <NUM_LIT:3> , "<STR_LIT>" <EOL> x , y , z = np . indices ( s . shape ) <EOL> elif len ( args ) == <NUM_LIT:3> : <EOL> x , y , z = args <EOL> s = None <EOL> elif len ( args ) == <NUM_LIT:4> : <EOL> x , y , z , s = args <EOL> if callable ( s ) : <EOL> s = s ( x , y , z ) <EOL> else : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> assert ( x . shape == y . shape and <EOL> y . shape == z . shape and <EOL> ( s is None <EOL> or s . shape == z . shape ) ) , "<STR_LIT>" <EOL> return x , y , z , s <EOL> def process_regular_2d_scalars ( * args , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> args = convert_to_arrays ( args ) <EOL> for index , arg in enumerate ( args ) : <EOL> if not callable ( arg ) : <EOL> args [ index ] = np . atleast_2d ( arg ) <EOL> if len ( args ) == <NUM_LIT:1> : <EOL> s = args [ <NUM_LIT:0> ] <EOL> assert len ( s . shape ) == <NUM_LIT:2> , "<STR_LIT>" <EOL> x , y = np . indices ( s . shape ) <EOL> elif len ( args ) == <NUM_LIT:3> : <EOL> x , y , s = args <EOL> if callable ( s ) : <EOL> s = s ( x , y ) <EOL> else : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> assert len ( s . shape ) == <NUM_LIT:2> , "<STR_LIT>" <EOL> if '<STR_LIT>' in kwargs : <EOL> mask = kwargs [ '<STR_LIT>' ] <EOL> s [ mask . astype ( '<STR_LIT:bool>' ) ] = np . nan <EOL> s = s . astype ( '<STR_LIT:float>' ) <EOL> return x , y , s <EOL> def vector_scatter ( * args , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> x , y , z , u , v , w = process_regular_vectors ( * args ) <EOL> scalars = kwargs . pop ( '<STR_LIT>' , None ) <EOL> if scalars is not None : <EOL> scalars = np . ravel ( scalars ) <EOL> name = kwargs . pop ( '<STR_LIT:name>' , '<STR_LIT>' ) <EOL> data_source = MGlyphSource ( ) <EOL> data_source . reset ( x = x , y = y , z = z , u = u , v = v , w = w , scalars = scalars ) <EOL> ds = tools . add_dataset ( data_source . dataset , name , ** kwargs ) <EOL> data_source . m_data = ds <EOL> return ds <EOL> def vector_field ( * args , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> if len ( args ) == <NUM_LIT:3> : <EOL> x = y = z = np . atleast_3d ( <NUM_LIT:1> ) <EOL> u , v , w = [ np . atleast_3d ( a ) for a in args ] <EOL> else : <EOL> x , y , z , u , v , w = [ np . atleast_3d ( a ) <EOL> for a in process_regular_vectors ( * args ) ] <EOL> scalars = kwargs . pop ( '<STR_LIT>' , None ) <EOL> if scalars is not None : <EOL> scalars = np . atleast_3d ( scalars ) <EOL> data_source = MArraySource ( ) <EOL> data_source . reset ( x = x , y = y , z = z , u = u , v = v , w = w , scalars = scalars ) <EOL> name = kwargs . pop ( '<STR_LIT:name>' , '<STR_LIT>' ) <EOL> return tools . add_dataset ( data_source . m_data , name , ** kwargs ) <EOL> def scalar_scatter ( * args , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> x , y , z , s = process_regular_scalars ( * args ) <EOL> if s is not None : <EOL> s = np . ravel ( s ) <EOL> data_source = MGlyphSource ( ) <EOL> data_source . reset ( x = x , y = y , z = z , scalars = s ) <EOL> name = kwargs . pop ( '<STR_LIT:name>' , '<STR_LIT>' ) <EOL> ds = tools . add_dataset ( data_source . dataset , name , ** kwargs ) <EOL> data_source . m_data = ds <EOL> return ds <EOL> def scalar_field ( * args , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> if len ( args ) == <NUM_LIT:1> : <EOL> x = y = z = np . atleast_1d ( <NUM_LIT:1> ) <EOL> s = args [ <NUM_LIT:0> ] <EOL> else : <EOL> x , y , z , s = process_regular_scalars ( * args ) <EOL> data_source = MArraySource ( ) <EOL> data_source . reset ( x = x , y = y , z = z , scalars = s ) <EOL> name = kwargs . pop ( '<STR_LIT:name>' , '<STR_LIT>' ) <EOL> return tools . add_dataset ( data_source . m_data , name , ** kwargs ) <EOL> def line_source ( * args , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> if len ( args ) == <NUM_LIT:1> : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> x , y , z , s = process_regular_scalars ( * args ) <EOL> data_source = MLineSource ( ) <EOL> data_source . reset ( x = x , y = y , z = z , scalars = s ) <EOL> name = kwargs . pop ( '<STR_LIT:name>' , '<STR_LIT>' ) <EOL> ds = tools . add_dataset ( data_source . dataset , name , ** kwargs ) <EOL> data_source . m_data = ds <EOL> return ds <EOL> def array2d_source ( * args , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> data_source = MArray2DSource ( ) <EOL> mask = kwargs . pop ( '<STR_LIT>' , None ) <EOL> if len ( args ) == <NUM_LIT:1> : <EOL> args = convert_to_arrays ( args ) <EOL> s = np . atleast_2d ( args [ <NUM_LIT:0> ] ) <EOL> data_source . reset ( scalars = s , mask = mask ) <EOL> else : <EOL> x , y , s = process_regular_2d_scalars ( * args , ** kwargs ) <EOL> data_source . reset ( x = x , y = y , scalars = s , mask = mask ) <EOL> name = kwargs . pop ( '<STR_LIT:name>' , '<STR_LIT>' ) <EOL> return tools . add_dataset ( data_source . m_data , name , ** kwargs ) <EOL> def grid_source ( x , y , z , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> scalars = kwargs . pop ( '<STR_LIT>' , None ) <EOL> if scalars is None : <EOL> scalars = z <EOL> mask = kwargs . pop ( '<STR_LIT>' , None ) <EOL> x , y , z , scalars = convert_to_arrays ( ( x , y , z , scalars ) ) <EOL> data_source = MGridSource ( ) <EOL> data_source . reset ( x = x , y = y , z = z , scalars = scalars , mask = mask ) <EOL> name = kwargs . pop ( '<STR_LIT:name>' , '<STR_LIT>' ) <EOL> ds = tools . add_dataset ( data_source . dataset , name , ** kwargs ) <EOL> data_source . m_data = ds <EOL> return ds <EOL> def vertical_vectors_source ( * args , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> if len ( args ) == <NUM_LIT:3> : <EOL> x , y , data = args <EOL> if np . isscalar ( x ) : <EOL> z = <NUM_LIT:0> <EOL> else : <EOL> z = np . zeros_like ( x ) <EOL> args = ( x , y , z , data ) <EOL> x , y , z , s = process_regular_scalars ( * args ) <EOL> if s is not None : <EOL> s = np . ravel ( s ) <EOL> data_source = MVerticalGlyphSource ( ) <EOL> data_source . reset ( x = x , y = y , z = z , scalars = s ) <EOL> name = kwargs . pop ( '<STR_LIT:name>' , '<STR_LIT>' ) <EOL> ds = tools . add_dataset ( data_source . dataset , name , ** kwargs ) <EOL> data_source . m_data = ds <EOL> return ds <EOL> def triangular_mesh_source ( x , y , z , triangles , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> x , y , z , triangles = convert_to_arrays ( ( x , y , z , triangles ) ) <EOL> if triangles . min ( ) < <NUM_LIT:0> : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> if triangles . max ( ) > x . size : <EOL> raise ValueError ( '<STR_LIT>' '<STR_LIT>' ) <EOL> scalars = kwargs . pop ( '<STR_LIT>' , None ) <EOL> if scalars is None : <EOL> scalars = z <EOL> data_source = MTriangularMeshSource ( ) <EOL> data_source . reset ( x = x , y = y , z = z , triangles = triangles , scalars = scalars ) <EOL> name = kwargs . pop ( '<STR_LIT:name>' , '<STR_LIT>' ) <EOL> ds = tools . add_dataset ( data_source . dataset , name , ** kwargs ) <EOL> data_source . m_data = ds <EOL> return ds <EOL> def open ( filename , figure = None ) : <EOL> """<STR_LIT>""" <EOL> if figure is None : <EOL> engine = tools . get_engine ( ) <EOL> elif figure is False : <EOL> engine = get_null_engine ( ) <EOL> else : <EOL> engine = engine_manager . find_figure_engine ( figure ) <EOL> engine . current_scene = figure <EOL> src = engine . open ( filename ) <EOL> return src <EOL> def _create_data_source ( metadata ) : <EOL> """<STR_LIT>""" <EOL> factory = metadata . get_callable ( ) <EOL> src = factory ( ) <EOL> engine = tools . get_engine ( ) <EOL> engine . add_source ( src ) <EOL> return src <EOL> def _make_functions ( namespace ) : <EOL> """<STR_LIT>""" <EOL> for src in registry . sources : <EOL> if len ( src . extensions ) == <NUM_LIT:0> : <EOL> func_name = camel2enthought ( src . id ) <EOL> if func_name . endswith ( '<STR_LIT>' ) : <EOL> func_name = func_name [ : - <NUM_LIT:7> ] <EOL> func = lambda metadata = src : _create_data_source ( metadata ) <EOL> func . __doc__ = src . help <EOL> func . __name__ = func_name <EOL> namespace [ func_name ] = func <EOL> __all__ . append ( func_name ) <EOL> _make_functions ( locals ( ) ) </s>
