Cleanup of widget invocation (#36)

- Remove passing redundant batch_size param
  - The buffer / batch size is available in trajectory
- Reset frame to the latest timestep before each widget invocation
  - Widget runs can iterate trajectory and hence required

Author: PardhavMaradani

mdadash/backend/analyses/dssp.py

@@ -154,7 +154,7 @@ def _compute_current_frame(self):
             self.u.trajectory.ts.data["time"],
         )
 
-    def _compute_batch(self, _batch_size):
+    def _compute_batch(self):
         """Compute values for current batch"""
         self.dssp.run()
         # `run()` can also be invoked with a non-serial backend
@@ -206,14 +206,14 @@ def run_every_frame(self):
         """every-frame run handler"""
         self._update_plot(self._compute_current_frame())
 
-    def run_batch(self, batch_size):
+    def run_batch(self):
         """batch run handler"""
-        self._update_plot(self._compute_batch(batch_size))
+        self._update_plot(self._compute_batch())
 
-    def get_parallel_job(self, batch_size):
+    def get_parallel_job(self):
         """get parallel job handler"""
         if self._run_frequency == "batch":
-            return delayed(self._compute_batch)(batch_size)
+            return delayed(self._compute_batch)()
         return delayed(self._compute_current_frame)()
 
     def apply_parallel_results(self, values):

mdadash/backend/analyses/rog.py

@@ -194,10 +194,10 @@ def _compute_current_frame(self):
             rog,
         )
 
-    def _compute_batch(self, batch_size):
+    def _compute_batch(self):
         """Compute ROG values for current batch"""
         values = []
-        for i in range(batch_size):
+        for i in range(self.u.trajectory.buffer_size):
             _ = self.u.trajectory[i]
             values.append(self._compute_current_frame())
         return values
@@ -225,14 +225,14 @@ def run_every_frame(self):
         """every-frame run handler"""
         self._update_plot(self._compute_current_frame())
 
-    def run_batch(self, batch_size):
+    def run_batch(self):
         """batch run handler"""
-        self._update_plot(self._compute_batch(batch_size))
+        self._update_plot(self._compute_batch())
 
-    def get_parallel_job(self, batch_size):
+    def get_parallel_job(self):
         """get parallel job handler"""
         if self._run_frequency == "batch":
-            return delayed(self._compute_batch)(batch_size)
+            return delayed(self._compute_batch)()
         return delayed(self._compute_current_frame)()
 
     def apply_parallel_results(self, values):

mdadash/backend/kernel/core.py

@@ -65,17 +65,21 @@ class BufferedTrajectory:
 
     """
 
-    def __init__(self, trajectory: mda.Universe.trajectory, batch_size: int):
+    def __init__(self, trajectory: mda.Universe.trajectory, buffer_size: int):
         self._trajectory = trajectory
-        self._batch_size = batch_size
-        self._buffer = deque(maxlen=batch_size)
+        self._buffer_size = buffer_size
+        self._buffer = deque(maxlen=buffer_size)
         self._buffer.append(trajectory.ts.copy())
         BufferedTrajectory.next.__doc__ = type(trajectory).next.__doc__
 
     @property
     def n_frames(self):
         return len(self._buffer)
 
+    @property
+    def buffer_size(self):
+        return self._buffer_size
+
     def __len__(self):
         return len(self._buffer)
 
@@ -389,7 +393,7 @@ async def _iter_loop(self):
                         batch_ready = (u.trajectory._frame + step) % (
                             step * batch_size
                         ) == 0
-                        self._wm.run_widgets(uid, batch_ready, batch_size)
+                        self._wm.run_widgets(uid, batch_ready)
                     # pylint: disable=broad-exception-caught
                     except Exception:  # pragma: no cover
                         logger.exception("Trajectory iteration failed for uid %d", uid)

mdadash/backend/tests/test_server.py

@@ -234,7 +234,7 @@ def run_every_frame(self):
     class _TestWidget4b(WidgetBase):
         name = "TestWidget4b"
 
-        def run_batch(self, batch_size):
+        def run_batch(self):
             pass
 
     # test duplicate widget name registraion exception

mdadash/backend/widgets/base.py

@@ -55,6 +55,10 @@ def _set_universe(self, u: mda.Universe):
         """Internal: Set the universe"""
         self.u = u
 
+    def _reset_frame_latest(self):
+        """Internal: Reset frame to latest timestep"""
+        _ = self.u.trajectory[-1]
+
     def _get_inputs(self):
         """Internal: Get the current instance inputs"""
         inputs = getattr(self, "_inputs")
@@ -172,33 +176,24 @@ def run_every_frame(self) -> None:
 
         """
 
-    def run_batch(self, batch_size: int) -> None:
+    def run_batch(self) -> None:
         """run_batch handler
 
         This handler is called every time a new batch of timesteps is full
         and ready to be run if the run frequency is set to `batch`
         (`_run_frequency='batch'`).
 
-        Parameters
-        ----------
-        batch_size: int
-            The batch size, which indicates the number of buffered timesteps
-            available in the buffer is passed
+        `self.u.trajectory.buffer_size` is the size of the buffer / batch
+        that can be used by the widget class.
 
         """
 
-    def get_parallel_job(self, batch_size: int) -> Any:
+    def get_parallel_job(self) -> Any:
         """get_parallel_job handler
 
         This handler is called if the run mode is set to `parallel`
         (`_run_mode='parallel'`) to get the parallel job to run.
 
-        Parameters
-        ----------
-        batch_size: int
-            The configured batch size to use in the parallel jobs
-            if their run frequency is `batch`
-
         Returns
         -------
         job: Any
@@ -272,7 +267,7 @@ def _validate_widget_class(cls, widget_class: WidgetBase) -> None:
         # check for one of the run methods to exist with correct params
         run_methods = {
             "run_every_frame": 1,
-            "run_batch": 2,
+            "run_batch": 1,
         }
         has_valid_run_method = False
         for run_method, n_params in run_methods.items():
@@ -546,11 +541,11 @@ def custom_setstate(self, state):
         IMDReader.__setstate__ = custom_setstate
         IMDReader.__getstate__ = custom_getstate
 
-    def _run_parallel_jobs(self, parallel_widgets, batch_size, parallel_results):
+    def _run_parallel_jobs(self, parallel_widgets, parallel_results):
         """Internal: Run parallel jobs using joblib.Parallel"""
         parallel_jobs = []
         for widget in parallel_widgets:
-            parallel_jobs.append(widget.get_parallel_job(batch_size))
+            parallel_jobs.append(widget.get_parallel_job())
         try:
             results = Parallel(
                 n_jobs=self.n_jobs, initializer=WidgetManager._patch_IMDReader
@@ -560,7 +555,7 @@ def _run_parallel_jobs(self, parallel_widgets, batch_size, parallel_results):
         except Exception:  # pragma: no cover
             logger.exception("Parallel run failed for jobs %s", parallel_jobs)
 
-    def run_widgets(self, uid: int, batch_ready: bool, batch_size: int) -> None:
+    def run_widgets(self, uid: int, batch_ready: bool) -> None:
         """Run widget instances
 
         Parameters
@@ -571,9 +566,6 @@ def run_widgets(self, uid: int, batch_ready: bool, batch_size: int) -> None:
         batch_ready: bool
             Flag indicating if a batch of timesteps is full
 
-        batch_size: int
-            Size of the batch
-
         """
         # collect widgets that need to be run
         parallel_widgets = []
@@ -594,19 +586,19 @@ def run_widgets(self, uid: int, batch_ready: bool, batch_size: int) -> None:
                 target=self._run_parallel_jobs,
                 args=(
                     parallel_widgets,
-                    batch_size,
                     parallel_results,
                 ),
             )
             parallel_thread.start()
         # run serial widgets
         for widget in serial_widgets:
+            widget._reset_frame_latest()
             with _widget_uuid_in_metadata(widget.uuid):
                 try:
                     if widget._run_frequency == "every-frame":
                         widget.run_every_frame()
                     elif batch_ready:
-                        widget.run_batch(batch_size)
+                        widget.run_batch()
                 # pylint: disable=broad-exception-caught
                 except Exception:  # pragma: no cover
                     logger.exception("Serial run failed for widget %s", widget.uuid)

mdadash/frontend/src/views/SettingsView.vue

@@ -163,12 +163,12 @@
 
               <v-number-input
                 class="mb-4"
-                label="Batch size"
+                label="Buffer / batch size"
                 variant="outlined"
                 v-model="settings.universe_configs[0].batch_size"
                 :min="1"
                 control-variant="hidden"
-                hint="Number of timesteps to batch for a batch run"
+                hint="Number of timesteps to buffer for a batch run"
                 persistent-hint
               ></v-number-input>