def __getitem__(self, value):
if isinstance(value, str):
print(list(self.__dict__.keys()))
try:
return self.__dict__['raw'][value]
except:
print('{} is not a valid data structure.'.format(value))
elif isinstance(value, int) or isinstance(value, np.int64):
# Otherwise, we assume they're trying to index for a seg_id.
#if value < self.parent.walkers:
current = {}
seg_items = ['weights', 'pcoord', 'auxdata', 'parents', 'seg_id', 'states']
current['pcoord'] = self.__dict__['raw']['pcoord'][value]
current['states'] = self.__dict__['raw']['states'][value]
current['bins'] = self.__dict__['raw']['bins'][value]
current['parents'] = self.__dict__['raw']['parents'][value]
current['seg_id'] = self.__dict__['raw']['seg_id'][value]
current['weights'] = self.__dict__['raw']['weights'][value]
try:
current['auxdata'] = {}
for key in list(self.__dict__['raw']['auxdata'].keys()):
current['auxdata'][key] = self.__dict__['raw']['auxdata'][key][value]
except:
pass
current = WIPIDataset(current, 'Segment {} in Iter {}'.format(value, self.iteration))
return current
def trace(self, seg_id):
'''
Runs a trace on a seg_id within the current iteration, all the way back to the beginning,
returning a dictionary containing all interesting information:
seg_id, pcoord, states, bins, weights, iteration, auxdata (optional)
sorted in chronological order.
Call with a seg_id.
'''
if seg_id >= self.current.walkers:
print("Walker seg_id # {} is beyond the max count of {} walkers.".format(seg_id, self.current.walkers))
return 1
pi = self.progress.indicator
with pi:
pi.new_operation('Tracing scheme:iter:seg_id {}:{}:{}'.format(self.scheme, self.iteration, seg_id), self.iteration)
current = { 'seg_id': [], 'pcoord': [], 'states': [], 'weights': [], 'iteration': [], 'bins': [] }
keys = []
try:
current['auxdata'] = {}
for key in list(self.current['auxdata'].keys()):
current['auxdata'][key] = []
key = []
except:
pass
for iter in reversed(list(range(1, self.iteration+1))):
iter_group = self.data_reader.get_iter_group(iter)
particles = self.data_reader.data_manager.get_iter_summary(int(iter))['n_particles']
current['pcoord'].append(iter_group['pcoord'][seg_id, :, :])
current['states'].append(self.assign['trajlabels'][iter-1, seg_id,:])
current['bins'].append(self.assign['assignments'][iter-1, seg_id,:])
current['seg_id'].append(seg_id)
current['weights'].append(iter_group['seg_index']['weight'][seg_id])
current['iteration'].append(iter)
try:
for key in keys:
current['auxdata'][key].append(iter_group['auxdata'][key][seg_id])
except:
pass
seg_id = iter_group['seg_index']['parent_id'][seg_id]
if seg_id < 0:
# Necessary for steady state simulations. This means they started in that iteration.
break
pi.progress += 1
current['seg_id'] = list(reversed(current['seg_id']))
current['iteration'] = list(reversed(current['iteration']))
current['states'] = np.concatenate(np.array(list(reversed(current['states']))))
current['bins'] = np.concatenate(np.array(list(reversed(current['bins']))))
current['weights'] = np.array(list(reversed(current['weights'])))
current['pcoord'] = np.concatenate(np.array(list(reversed(current['pcoord']))))
try:
for key in keys():
current['auxdata'][key] = np.concatenate(np.array(list(reversed(current['auxdata'][key]))))
except:
pass
current['state_labels'] = self.assign['state_labels']
for i in ['pcoord', 'states', 'bins', 'weights']:
current[i] = WIPIDataset(raw=current[i], key=i)
if i == 'weights':
current[i].plotter = Plotter(np.log10(current[i].raw), str('log10 of ' + str(i)), iteration=current[i].raw.shape[0], interface=self.interface)
else:
current[i].plotter = Plotter(current[i].raw, i, iteration=current[i].raw.shape[0], interface=self.interface)
current[i].plot = current[i].plotter.plot
return WIPIDataset(raw=current, key=seg_id)