Add option to save timestamp positions [WIP]

Author: tritemio

notebooks/PyBroMo - 1. Simulate 3D trajectories - single core.ipynb

@@ -95,8 +95,8 @@
     "psf = pbm.NumericPSF()\n",
     "\n",
     "# Particles definition\n",
-    "P = pbm.Particles(num_particles=20, D=D1, box=box, rs=rs)\n",
-    "P.add(num_particles=15, D=D2)\n",
+    "P = pbm.Particles(num_particles=2, D=D1, box=box, rs=rs)\n",
+    "P.add(num_particles=2, D=D2)\n",
     "\n",
     "# Simulation time step (seconds)\n",
     "t_step = 0.5e-6\n",
@@ -214,7 +214,8 @@
    },
    "outputs": [],
    "source": [
-    "S.simulate_diffusion(total_emission=False, save_pos=True, verbose=True, rs=rs, chunksize=2**19, chunkslice='times')"
+    "S.simulate_diffusion(total_emission=False, save_pos=True, verbose=True, \n",
+    "                     rs=rs, chunksize=2**19, chunkslice='times')"
    ]
   },
   {
@@ -228,6 +229,16 @@
     "print('Current random state:', pbm.hash_(rs.get_state()))"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "position = S.position[:]\n",
+    "position.shape, position.dtype"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -273,7 +284,7 @@
    },
    "outputs": [],
    "source": [
-    "S = pbm.ParticlesSimulation.from_datafile('0168')  # Read-only by default"
+    "S = pbm.ParticlesSimulation.from_datafile('9125')  # Read-only by default"
    ]
   },
   {
@@ -352,7 +363,6 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "collapsed": true,
     "scrolled": true
    },
    "outputs": [],
@@ -415,9 +425,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "collapsed": true
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "S.store.close()"
@@ -440,7 +448,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "S = pbm.ParticlesSimulation.from_datafile('0168', mode='w')"
+    "S = pbm.ParticlesSimulation.from_datafile('9125', mode='w')"
    ]
   },
   {
@@ -459,7 +467,7 @@
    "outputs": [],
    "source": [
     "S.simulate_timestamps_mix(max_rates=(200e3,), \n",
-    "                          populations=(slice(0, 35),),\n",
+    "                          populations=(slice(None),),\n",
     "                          bg_rate=1e3)"
    ]
   },
@@ -547,19 +555,17 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "collapsed": true
-   },
+   "metadata": {},
    "outputs": [],
    "source": []
   }
  ],
  "metadata": {
   "anaconda-cloud": {},
   "kernelspec": {
-   "display_name": "Python 3.6 (pybromo_env)",
+   "display_name": "Python 3.6 (py36)",
    "language": "python",
-   "name": "pybromo_env"
+   "name": "py36"
   },
   "language_info": {
    "codemirror_mode": {
@@ -571,7 +577,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.6.3"
+   "version": "3.6.7"
   },
   "toc": {
    "nav_menu": {},

pybromo/diffusion.py

@@ -686,7 +686,7 @@ def timestamp_names(self):
         return names
 
     def _sim_timestamps(self, max_rate, bg_rate, emission, i_start, rs,
-                        ip_start=0, scale=10, sort=True):
+                        ip_start=0, scale=10, sort=True, save_pos=False):
         """Simulate timestamps from emission trajectories.
 
         Uses attributes: `.t_step`.
@@ -702,7 +702,9 @@ def _sim_timestamps(self, max_rate, bg_rate, emission, i_start, rs,
         assert counts_chunk.shape == (nrows, emission.shape[1])
         max_counts = counts_chunk.max()
         if max_counts == 0:
-            return np.array([], dtype=np.int64), np.array([], dtype=np.int64)
+            return (np.array([], dtype=np.int64),    # timestamps
+                    np.array([], dtype=np.int64),    # particles
+                    np.array([], dtype=np.float32))  # positions
 
         time_start = i_start * scale
         time_stop = time_start + counts_chunk.shape[1] * scale
@@ -722,51 +724,67 @@ def _sim_timestamps(self, max_rate, bg_rate, emission, i_start, rs,
             # Append current particle
             times_chunk_p.append(t)
             par_index_chunk_p.append(np.full(t.size, ip + ip_start, dtype='u1'))
+            if save_pos:
+                # TODO compute and asve positions
+                pass
 
         # Merge the arrays of different particles
         times_chunk = np.hstack(times_chunk_p)
         par_index_chunk = np.hstack(par_index_chunk_p)
+        positions_chunk = np.zeros(
+            (times_chunk.shape[-1], self.positions.shape[1]), dtype=np.float32)
 
         if sort:
             # Sort timestamps inside the merged chunk
             index_sort = times_chunk.argsort(kind='mergesort')
             times_chunk = times_chunk[index_sort]
             par_index_chunk = par_index_chunk[index_sort]
+            if save_pos:
+                # TODO: reorder positions
+                pass
 
-        return times_chunk, par_index_chunk
+        return times_chunk, par_index_chunk, positions_chunk
 
     def _sim_timestamps_populations(self, emission, max_rates, populations,
-                                    bg_rates, i_start, rs, scale=10):
+                                    bg_rates, i_start, rs, scale=10,
+                                    save_pos=False):
         if populations is None:
             populations = [slice(0, self.num_particles)]
 
         # Loop for each population
         ts_chunk_pop_list, par_index_chunk_pop_list = [], []
+        positions_chunk_pop_list = []
         # Loop through populations
         for rate, pop, bg in zip(max_rates, populations, bg_rates):
             emission_pop = emission[pop]
-            ts_chunk_pop, par_index_chunk_pop = \
+            ts_chunk_pop, par_index_chunk_pop, positions_pop = \
                 self._sim_timestamps(
                     rate, bg, emission_pop, i_start, ip_start=pop.start,
-                    rs=rs, scale=scale, sort=False)
+                    rs=rs, scale=scale, sort=False, save_pos=save_pos)
 
             ts_chunk_pop_list.append(ts_chunk_pop)
             par_index_chunk_pop_list.append(par_index_chunk_pop)
+            if save_pos:
+                positions_chunk_pop_list.append(positions_pop)
 
         # Merge populations
         times_chunk_s = np.hstack(ts_chunk_pop_list)
         par_index_chunk_s = np.hstack(par_index_chunk_pop_list)
+        if save_pos:
+            positions_chunk_s = np.hstack(positions_chunk_pop_list)
 
         # Sort timestamps inside the merged chunk
         index_sort = times_chunk_s.argsort(kind='mergesort')
         times_chunk_s = times_chunk_s[index_sort]
         par_index_chunk_s = par_index_chunk_s[index_sort]
-        return times_chunk_s, par_index_chunk_s
+        if save_pos:
+            positions_chunk_s = positions_chunk_s[index_sort]
+        return times_chunk_s, par_index_chunk_s, positions_chunk_s
 
     def simulate_timestamps_mix(self, max_rates, populations, bg_rate,
                                 rs=None, seed=1, chunksize=2**16,
                                 comp_filter=None, overwrite=False,
-                                skip_existing=False, scale=10,
+                                skip_existing=False, scale=10, save_pos=False,
                                 path=None, t_chunksize=None, timeslice=None):
         """Compute one timestamps array for a mixture of N populations.
 
@@ -797,6 +815,8 @@ def simulate_timestamps_mix(self, max_rates, populations, bg_rate,
             scale (int): `self.t_step` is multiplied by `scale` to obtain the
                 timestamps units in seconds.
             path (string): folder where to save the data.
+            save_pos (bool): if True save 3D position of each particle
+                for each emitted photon.
             timeslice (float or None): timestamps are simulated until
                 `timeslice` seconds. If None, simulate until `self.t_max`.
         """
@@ -814,13 +834,16 @@ def simulate_timestamps_mix(self, max_rates, populations, bg_rate,
             max_rates=max_rates, bg_rate=bg_rate, populations=populations,
             num_particles=self.num_particles,
             bg_particle=self.num_particles,
-            overwrite=overwrite, chunksize=chunksize
+            overwrite=overwrite, chunksize=chunksize,
+            save_pos=save_pos,
             )
+        if save_pos:
+            kw.update(spatial_dims=self.positions.shape[1])
         if comp_filter is not None:
             kw.update(comp_filter=comp_filter)
         try:
-            self._timestamps, self._tparticles = (self.ts_store
-                                                  .add_timestamps(**kw))
+            self._timestamps, self._tparticles, self._tpositions = (
+                self.ts_store.add_timestamps(**kw))
         except ExistingArrayError as e:
             if skip_existing:
                 print(' - Skipping already present timestamps array.')
@@ -832,7 +855,7 @@ def simulate_timestamps_mix(self, max_rates, populations, bg_rate,
         self._timestamps.attrs['init_random_state'] = rs.get_state()
         self._timestamps.attrs['PyBroMo'] = __version__
 
-        ts_list, part_list = [], []
+        ts_list, part_list, pos_list = [], [], []
         # Load emission in chunks, and save only the final timestamps
         bg_rates = [None] * (len(max_rates) - 1) + [bg_rate]
         prev_time = 0
@@ -846,18 +869,21 @@ def simulate_timestamps_mix(self, max_rates, populations, bg_rate,
 
             em_chunk = self.emission[:, i_start:i_end]
 
-            times_chunk_s, par_index_chunk_s = \
+            times_chunk_s, par_index_chunk_s, positions_chunk_s = \
                 self._sim_timestamps_populations(
                     em_chunk, max_rates, populations, bg_rates, i_start,
-                    rs, scale)
+                    rs, scale, save_pos=save_pos)
 
-            # Save sorted timestamps (suffix '_s') and corresponding particles
+            # Save sorted "photons" (suffix '_s')
             ts_list.append(times_chunk_s)
             part_list.append(par_index_chunk_s)
+            pos_list.append(positions_chunk_s)  # it may be a list of None
 
-        for ts, part in zip(ts_list, part_list):
+        for ts, part, pos in zip(ts_list, part_list, pos_list):
             self._timestamps.append(ts)
             self._tparticles.append(part)
+            if save_pos:
+                self._tpositions.append(pos)
 
         # Save current random state so it can be resumed in the next session
         self.ts_group._v_attrs['last_random_state'] = rs.get_state()
@@ -933,8 +959,8 @@ def simulate_timestamps_mix_da(self, max_rates_d, max_rates_a,
 
         kw.update(name=name_d, max_rates=max_rates_d, bg_rate=bg_rate_d)
         try:
-            self._timestamps_d, self._tparticles_d = (self.ts_store
-                                                      .add_timestamps(**kw))
+            self._timestamps_d, self._tparticles_d, _ = (
+                self.ts_store.add_timestamps(**kw))
         except ExistingArrayError as e:
             if skip_existing:
                 print(' - Skipping already present timestamps array.')
@@ -944,8 +970,8 @@ def simulate_timestamps_mix_da(self, max_rates_d, max_rates_a,
 
         kw.update(name=name_a, max_rates=max_rates_a, bg_rate=bg_rate_a)
         try:
-            self._timestamps_a, self._tparticles_a = (self.ts_store
-                                                      .add_timestamps(**kw))
+            self._timestamps_a, self._tparticles_a, _ = (
+                self.ts_store.add_timestamps(**kw))
         except ExistingArrayError as e:
             if skip_existing:
                 print(' - Skipping already present timestamps array.')
@@ -972,12 +998,12 @@ def simulate_timestamps_mix_da(self, max_rates_d, max_rates_a,
 
             em_chunk = self.emission[:, i_start:i_end]
 
-            times_chunk_s_d, par_index_chunk_s_d = \
+            times_chunk_s_d, par_index_chunk_s_d, _ = \
                 self._sim_timestamps_populations(
                     em_chunk, max_rates_d, populations, bg_rates_d, i_start,
                     rs, scale)
 
-            times_chunk_s_a, par_index_chunk_s_a = \
+            times_chunk_s_a, par_index_chunk_s_a, _ = \
                 self._sim_timestamps_populations(
                     em_chunk, max_rates_a, populations, bg_rates_a, i_start,
                     rs, scale)
@@ -1057,8 +1083,8 @@ def simulate_timestamps_mix_da_online(self, max_rates_d, max_rates_a,
 
         kw.update(name=name_d, max_rates=max_rates_d, bg_rate=bg_rate_d)
         try:
-            self._timestamps_d, self._tparticles_d = (self.ts_store
-                                                      .add_timestamps(**kw))
+            self._timestamps_d, self._tparticles_d, _ = (
+                self.ts_store.add_timestamps(**kw))
         except ExistingArrayError as e:
             if skip_existing:
                 print(' - Skipping, timestamps array already present.')
@@ -1068,8 +1094,8 @@ def simulate_timestamps_mix_da_online(self, max_rates_d, max_rates_a,
 
         kw.update(name=name_a, max_rates=max_rates_a, bg_rate=bg_rate_a)
         try:
-            self._timestamps_a, self._tparticles_a = (self.ts_store
-                                                      .add_timestamps(**kw))
+            self._timestamps_a, self._tparticles_a, _ = (
+                self.ts_store.add_timestamps(**kw))
         except ExistingArrayError as e:
             if skip_existing:
                 print(' - Skipping, timestamps array already present.')
@@ -1103,12 +1129,12 @@ def simulate_timestamps_mix_da_online(self, max_rates_d, max_rates_a,
                                                  save_pos=False, radial=False,
                                                  wrap_func=wrap_periodic)
 
-            times_chunk_s_d, par_index_chunk_s_d = \
+            times_chunk_s_d, par_index_chunk_s_d, _ = \
                 self._sim_timestamps_populations(
                     em_chunk, max_rates_d, populations, bg_rates_d, i_start,
                     rs, scale)
 
-            times_chunk_s_a, par_index_chunk_s_a = \
+            times_chunk_s_a, par_index_chunk_s_a, _ = \
                 self._sim_timestamps_populations(
                     em_chunk, max_rates_a, populations, bg_rates_a, i_start,
                     rs, scale)

pybromo/storage.py

@@ -255,7 +255,8 @@ def __init__(self, datafile, path='./', nparams=dict(), attr_params=dict(),
     def add_timestamps(self, name, clk_p, max_rates, bg_rate,
                        num_particles, bg_particle, populations=None,
                        overwrite=False, chunksize=2**16,
-                       comp_filter=default_compression):
+                       comp_filter=default_compression, save_pos=False,
+                       spatial_dims=None):
         if name in self.h5file.root.timestamps:
             if overwrite:
                 self.h5file.remove_node('/timestamps', name=name)
@@ -286,7 +287,18 @@ def add_timestamps(self, name, clk_p, max_rates, bg_rate,
         particles_array.set_attr('bg_particle', bg_particle)
         particles_array.set_attr('PyBroMo', __version__)
         particles_array.set_attr('creation_time', current_time())
-        return times_array, particles_array
+        positions_array = None
+        if save_pos:
+            assert spatial_dims is not None, 'You need to pass `spatial_dims`.'
+            positions_array = self.h5file.create_earray(
+                '/timestamps', name + '_pos', atom=tables.Float32Atom(),
+                shape=(0, spatial_dims),
+                chunkshape=(chunksize,),
+                filters=comp_filter,
+                title='Particle position for each timestamp')
+            positions_array.set_attr('PyBroMo', __version__)
+            positions_array.set_attr('creation_time', current_time())
+        return times_array, particles_array, positions_array
 
 
 if __name__ == '__main__':