Merge pull request #217 from prabhuramachandran/more-viewer-improvements

More viewer improvements

pysph/tools/mayavi_viewer.py

@@ -12,6 +12,7 @@
 import numpy
 import os
 import os.path
+import time
 
 if not os.environ.get('ETS_TOOLKIT'):
     # Set the default toolkit to qt4 unless the user has explicitly
@@ -28,7 +29,7 @@
 from mayavi.core.api import PipelineBase  # noqa: E402
 from mayavi.core.ui.api import (
     MayaviScene, SceneEditor, MlabSceneModel)  # noqa: E402
-from pyface.timer.api import Timer, do_later  # noqa: E402
+from pyface.timer.api import Timer, do_later, do_after  # noqa: E402
 from tvtk.api import tvtk  # noqa: E402
 from tvtk.array_handler import array2vtk  # noqa: E402
 
@@ -64,6 +65,16 @@ def set_arrays(dataset, particle_array):
     dataset._update_data()
 
 
+def get_files_in_dir(pth):
+    '''Get the files in a given directory.
+    '''
+    _files = glob.glob(os.path.join(pth, '*.hdf5'))
+    if len(_files) == 0:
+        _files = glob.glob(os.path.join(pth, '*.npz'))
+        _files = [x for x in _files if os.path.basename(x) != 'results.npz']
+    return _files
+
+
 def glob_files(fname):
     """Glob for all similar files given one of them.
 
@@ -340,6 +351,7 @@ class ParticleArrayHelper(HasTraits):
                     Item(name='visible'),
                     Item(name='show_legend', label='Legend'),
                     Item(name='scalar',
+                         enabled_when='len(formula) == 0',
                          editor=EnumEditor(name='scalar_list')),
                     Item(name='list_all_scalars', label='All scalars'),
                     Item(name='show_time', label='Time'),
@@ -644,9 +656,12 @@ class MayaviViewer(HasTraits):
                                  'or from saved files')
 
     shell = Button('Launch Python Shell')
-    host = Str('localhost', desc='machine to connect to')
-    port = Int(8800, desc='port to use to connect to solver')
-    authkey = Password('pysph', desc='authorization key')
+    host = Str('localhost', enter_set=True, auto_set=False,
+               desc='machine to connect to')
+    port = Int(8800, enter_set=True, auto_set=False,
+               desc='port to use to connect to solver')
+    authkey = Password('pysph', enter_set=True, auto_set=False,
+                       desc='authorization key')
     host_changed = Bool(True)
     client = Instance(MultiprocessingClient)
     controller = Property(depends_on='live_mode, host_changed')
@@ -669,8 +684,10 @@ class MayaviViewer(HasTraits):
     ########################################
     # Timer traits.
     timer = Instance(Timer)
-    interval = Float(5.0, enter_set=True, auto_set=False,
-                     desc='frequency in seconds with which plot is updated')
+    interval = Float(
+        5.0, enter_set=True, auto_set=False,
+        desc='suggested frequency in seconds with which plot is updated'
+    )
 
     ########################################
     # Solver info/control.
@@ -691,6 +708,8 @@ class MayaviViewer(HasTraits):
     _solver_data = Any
     _file_name = Str
     _particle_array_updated = Bool
+    _doing_update = Bool(False)
+    _poll_interval = Float(5.0)
 
     ########################################
     # The layout of the dialog created
@@ -800,13 +819,13 @@ def start_timer(self):
         # Just accessing the timer will start it.
         t = self.timer
         if not t.IsRunning():
-            t.Start(int(self.interval*1000))
+            t.Start(int(self._poll_interval*1000))
 
     @on_trait_change('scene:activated')
     def update_plot(self):
 
         # No need to do this if files are being used.
-        if not self.live_mode:
+        if self._doing_update or not self.live_mode:
             return
 
         # do not update if solver is paused
@@ -820,25 +839,39 @@ def update_plot(self):
         if controller is None:
             return
 
-        self.current_time = t = controller.get_t()
-        self.time_step = controller.get_dt()
-        self.iteration = controller.get_count()
-
-        arrays = []
-        for idx, name in enumerate(self.pa_names):
-            pa = controller.get_named_particle_array(name)
-            arrays.append(pa)
-            pah = self.particle_arrays[idx]
-            pah.set(particle_array=pa, time=t)
+        try:
+            start = time.time()
+            self._doing_update = True
+            self.current_time = t = controller.get_t()
+            self.time_step = controller.get_dt()
+            self.iteration = controller.get_count()
+
+            arrays = []
+            for idx, name in enumerate(self.pa_names):
+                pa = controller.get_named_particle_array(name)
+                arrays.append(pa)
+                pah = self.particle_arrays[idx]
+                pah.set(particle_array=pa, time=t)
 
-        self.interpolator.particle_arrays = arrays
+            self.interpolator.particle_arrays = arrays
 
-        if self.record:
-            self._do_snap()
+            total = time.time() - start
+            if total*3 > self._poll_interval or total*5 < self._poll_interval:
+                self._poll_interval = max(3*total, self.interval)
+                self._interval_changed(self._poll_interval)
+            if self.record:
+                self._do_snap()
+        finally:
+            self._doing_update = False
 
     def run_script(self, path):
         """Execute a script in the namespace of the viewer.
         """
+        pas = self.particle_arrays
+        if len(pas) == 0 or pas[0].plot is None:
+            do_after(2000, self.run_script, path)
+            return
+
         with open(path) as fp:
             data = fp.read()
             ns = self._get_shell_namespace()
@@ -942,6 +975,7 @@ def _client_changed(self, old, new):
             for x in self.pa_names
         ]
         self.interpolator = InterpolatorView(scene=self.scene)
+        do_later(self.update_plot)
 
         output_dir = self.client.controller.get_output_directory()
         config_file = os.path.join(output_dir, 'mayavi_config.py')
@@ -965,10 +999,11 @@ def _interval_changed(self, value):
             return
         if t.IsRunning():
             t.Stop()
-            t.Start(int(value*1000))
+            interval = max(value, self._poll_interval)
+            t.Start(int(interval*1000))
 
     def _timer_default(self):
-        return Timer(int(self.interval*1000), self._timer_event)
+        return Timer(int(self._poll_interval*1000), self._timer_event)
 
     def _pause_solver_changed(self, value):
         if self.live_mode:
@@ -1007,7 +1042,7 @@ def _files_changed(self, value):
         else:
             if not t.IsRunning():
                 t.Stop()
-                t.Start(self.interval*1000)
+                t.Start(self._poll_interval*1000)
 
     def _file_count_changed(self, value):
         # Save out any updates for the previous file if needed.
@@ -1112,8 +1147,7 @@ def _get_shell_namespace(self):
                     mlab=self.scene.mlab)
 
     def _directory_changed(self, d):
-        ext = os.path.splitext(self.files[-1])[1]
-        files = glob.glob(os.path.join(d, '*' + ext))
+        files = get_files_in_dir(d)
         if len(files) > 0:
             self._clear()
             sort_file_list(files)
@@ -1264,9 +1298,7 @@ def main(args=None):
                 files.extend(glob.glob(arg))
                 continue
             elif os.path.isdir(arg):
-                _files = glob.glob(os.path.join(arg, '*.hdf5'))
-                if len(_files) == 0:
-                    _files = glob.glob(os.path.join(arg, '*.npz'))
+                _files = get_files_in_dir(arg)
                 files.extend(_files)
                 config_file = os.path.join(arg, 'mayavi_config.py')
                 if os.path.exists(config_file):
@@ -1279,7 +1311,7 @@ def main(args=None):
         try:
             val = eval(arg, math.__dict__)
             # this will fail if arg is a string.
-        except NameError:
+        except Exception:
             val = arg
         kw[key] = val