class MCMC(Sampler): """ This class fits probability models using Markov Chain Monte Carlo. Each stochastic variable is assigned a StepMethod object, which makes it take a single MCMC step conditional on the rest of the model. These step methods are called in turn. >>> A = MCMC(input, db, verbose=0) :Parameters: - input : module, list, tuple, dictionary, set, object or nothing. Model definition, in terms of Stochastics, Deterministics, Potentials and Containers. If nothing, all nodes are collected from the base namespace. - db : string The name of the database backend that will store the values of the stochastics and deterministics sampled during the MCMC loop. - verbose : integer Level of output verbosity: 0=none, 1=low, 2=medium, 3=high Inherits all methods and attributes from Model. Subclasses must define the _loop method: - _loop(self, *args, **kwargs): Can be called after a sampling run is interrupted (by pausing, halting or a KeyboardInterrupt) to continue the sampling run. _loop must be able to handle KeyboardInterrupts gracefully, and should monitor the sampler's status periodically. Available status values are: - 'ready': Ready to sample. - 'paused': A pause has been requested, or the sampler is paused. _loop should return control as soon as it is safe to do so. - 'halt': A halt has been requested, or the sampler is stopped. _loop should call halt as soon as it is safe to do so. - 'running': Sampling is in progress. :SeeAlso: Model, Sampler, StepMethod. """ def __init__(self, input=None, db='ram', name='MCMC', calc_deviance=True, **kwds): """Initialize an MCMC instance. :Parameters: - input : module, list, tuple, dictionary, set, object or nothing. Model definition, in terms of Stochastics, Deterministics, Potentials and Containers. If nothing, all nodes are collected from the base namespace. - db : string The name of the database backend that will store the values of the stochastics and deterministics sampled during the MCMC loop. - verbose : integer Level of output verbosity: 0=none, 1=low, 2=medium, 3=high - **kwds : Keywords arguments to be passed to the database instantiation method. """ Sampler.__init__(self, input, db, name, calc_deviance=calc_deviance, **kwds) self._sm_assigned = False self.step_method_dict = {} for s in self.stochastics: self.step_method_dict[s] = [] self._state = ['status', '_current_iter', '_iter', '_tune_interval', '_burn', '_thin', '_tuned_count', '_tune_throughout', '_burn_till_tuned'] def use_step_method(self, step_method_class, *args, **kwds): """ M.use_step_method(step_method_class, *args, **kwds) Example of usage: To handle stochastic A with a Metropolis instance, M.use_step_method(Metropolis, A, sig=.1) To subsequently get a reference to the new step method, S = M.step_method_dict[A][0] """ new_method = step_method_class(*args, **kwds) if self.verbose > 1: print_('Using step method %s. Stochastics: ' % step_method_class.__name__) for s in new_method.stochastics: self.step_method_dict[s].append(new_method) if self.verbose > 1: print_('\t'+s.__name__) if self._sm_assigned: self.step_methods.add(new_method) setattr(new_method, '_model', self) def remove_step_method(self, step_method): """ Removes a step method. """ try: for s in step_method.stochastics: self.step_method_dict[s].remove(step_method) if hasattr(self, "step_methods"): self.step_methods.discard(step_method) self._sm_assigned = False except AttributeError: for sm in step_method: self.remove_step_method(sm) def assign_step_methods(self, verbose=-1, draw_from_prior_when_possible = True): """ Make sure every stochastic variable has a step method. If not, assign a step method from the registry. """ if not self._sm_assigned: if draw_from_prior_when_possible: # Assign dataless stepper first last_gen = set([]) for s in self.stochastics - self.observed_stochastics: if s._random is not None: if len(s.extended_children)==0: last_gen.add(s) dataless, dataless_gens = crawl_dataless(set(last_gen), [last_gen]) if len(dataless): new_method = DrawFromPrior(dataless, dataless_gens[::-1], verbose=verbose) setattr(new_method, '_model', self) for d in dataless: if not d.observed: self.step_method_dict[d].append(new_method) if self.verbose > 1: print_('Assigning step method %s to stochastic %s' % (new_method.__class__.__name__, d.__name__)) for s in self.stochastics: # If not handled by any step method, make it a new step method using the registry if len(self.step_method_dict[s])==0: new_method = assign_method(s, verbose=verbose) setattr(new_method, '_model', self) self.step_method_dict[s].append(new_method) if self.verbose > 1: print_('Assigning step method %s to stochastic %s' % (new_method.__class__.__name__, s.__name__)) self.step_methods = set() for s in self.stochastics: self.step_methods |= set(self.step_method_dict[s]) for sm in self.step_methods: if sm.tally: for name in sm._tuning_info: self._funs_to_tally[sm._id+'_'+name] = lambda name=name, sm=sm: getattr(sm, name) else: # Change verbosity for step methods for sm_key in self.step_method_dict: for sm in self.step_method_dict[sm_key]: sm.verbose = verbose self.restore_sm_state() self._sm_assigned = True def sample(self, iter, burn=0, thin=1, tune_interval=1000, tune_throughout=True, save_interval=None, burn_till_tuned=False, stop_tuning_after=5, verbose=0, progress_bar=True): """ sample(iter, burn, thin, tune_interval, tune_throughout, save_interval, verbose, progress_bar) Initialize traces, run sampling loop, clean up afterward. Calls _loop. :Parameters: - iter : int Total number of iterations to do - burn : int Variables will not be tallied until this many iterations are complete, default 0 - thin : int Variables will be tallied at intervals of this many iterations, default 1 - tune_interval : int Step methods will be tuned at intervals of this many iterations, default 1000 - tune_throughout : boolean If true, tuning will continue after the burnin period (True); otherwise tuning will halt at the end of the burnin period. - save_interval : int or None If given, the model state will be saved at intervals of this many iterations - verbose : boolean - progress_bar : boolean Display progress bar while sampling. - burn_till_tuned: boolean If True the Sampler would burn samples until all step methods are tuned. A tuned step methods is one that was not tuned for the last `stop_tuning_after` tuning intervals. The burn-in phase will have a minimum of 'burn' iterations but could be longer if tuning is needed. After the phase is done the sampler will run for another (iter - burn) iterations, and will tally the samples according to the 'thin' argument. This means that the total number of iteration is update throughout the sampling procedure. If burn_till_tuned is True it also overrides the tune_thorughout argument, so no step method will be tuned when sample are being tallied. - stop_tuning_after: int the number of untuned successive tuning interval needed to be reach in order for the burn-in phase to be done (If burn_till_tuned is True). """ self.assign_step_methods(verbose=verbose) if burn > iter: raise ValueError('Burn interval cannot be larger than specified number of iterations.') self._n_tally = int(iter) - int(burn) if burn_till_tuned: self._stop_tuning_after = stop_tuning_after tune_throughout = False if verbose > 0: print "burn_til_tuned is True. tune_throughout is set to False" burn = int(max(burn, stop_tuning_after * tune_interval)) iter = self._n_tally + burn self._iter = int(iter) self._burn = int(burn) self._thin = int(thin) self._tune_interval = int(tune_interval) self._tune_throughout = tune_throughout self._burn_till_tuned = burn_till_tuned self._save_interval = save_interval length = max(int(np.floor((1.0*iter-burn)/thin)), 1) self.max_trace_length = length # Flags for tuning self._tuning = True self._tuned_count = 0 # Progress bar self.pbar = None if not verbose and progress_bar: self.pbar = ProgressBar(self._iter) # Run sampler Sampler.sample(self, iter, length, verbose) def _loop(self): # Set status flag self.status='running' try: while self._current_iter < self._iter and not self.status == 'halt': if self.status == 'paused': break i = self._current_iter # Update progress bar if not (i+1) % 100 and self.pbar: self.pbar.animate(i+1) # Tune at interval if i and not (i % self._tune_interval) and self._tuning: self.tune() #update _burn and _iter if needed if self._burn_till_tuned and (self._tuned_count == 0): new_burn = self._current_iter + int(self._stop_tuning_after * self._tune_interval) self._burn = max(new_burn, self._burn); self._iter = self._burn + self._n_tally self.pbar = ProgressBar(self._iter) self.pbar.animate(i+1) # Manage burn-in if i == self._burn: if self.verbose>0: print_('\nBurn-in interval complete') if not self._tune_throughout: self._tuning = False # Tell all the step methods to take a step for step_method in self.step_methods: if self.verbose > 2: print_('Step method %s stepping' % step_method._id) # Step the step method step_method.step() # Record sample to trace, if appropriate if i % self._thin == 0 and i >= self._burn: self.tally() if self._save_interval is not None: if i % self._save_interval==0: self.save_state() # Periodically commit samples to backend if not i % 1000: self.commit() # Increment interation self._current_iter += 1 except KeyboardInterrupt: self.status='halt' finally: # Stop progress bar if self.pbar: self.pbar.animate(self._iter) if self.status == 'halt': self._halt() def tune(self): """ Tell all step methods to tune themselves. """ # ======================================= # = This is what makes price.py puke... = # ======================================= # Only tune during burn-in # if self._current_iter > self._burn: # self._tuning = False # return if self.verbose > 0: print_('\tTuning at iteration', self._current_iter) # Initialize counter for number of tuning stochastics tuning_count = 0 for step_method in self.step_methods: verbose = self.verbose if step_method.verbose > -1: verbose = step_method.verbose # Tune step methods tuning_count += step_method.tune(verbose=self.verbose) if verbose > 1: print_('\t\tTuning step method %s, returned %i\n' %(step_method._id, tuning_count)) sys.stdout.flush() if self._burn_till_tuned: if not tuning_count: # If no step methods needed tuning, increment count self._tuned_count += 1 else: # Otherwise re-initialize count self._tuned_count = 0 # n consecutive clean intervals removed tuning # n is equal to self._stop_tuning_after if self._tuned_count == self._stop_tuning_after: if self.verbose > 0: print_('\nFinished tuning') self._tuning = False def get_state(self): """ Return the sampler and step methods current state in order to restart sampling at a later time. """ self.step_methods = set() for s in self.stochastics: self.step_methods |= set(self.step_method_dict[s]) state = Sampler.get_state(self) state['step_methods'] = {} # The state of each StepMethod. for sm in self.step_methods: state['step_methods'][sm._id] = sm.current_state().copy() return state def restore_sm_state(self): sm_state = self.db.getstate() if sm_state is not None: sm_state = sm_state.get('step_methods', {}) # Restore stepping methods state for sm in self.step_methods: sm.__dict__.update(sm_state.get(sm._id, {})) def _calc_dic(self): """Calculates deviance information Criterion""" # Find mean deviance mean_deviance = np.mean(self.db.trace('deviance')(), axis=0) # Set values of all parameters to their mean for stochastic in self.stochastics: # Calculate mean of paramter try: mean_value = np.mean(self.db.trace(stochastic.__name__)(), axis=0) # Set current value to mean stochastic.value = mean_value except KeyError: print_("No trace available for %s. DIC value may not be valid." % stochastic.__name__) # Return twice deviance minus deviance at means return 2*mean_deviance - self.deviance # Make dic a property def _get_dic(self): return self._calc_dic() dic = property(_get_dic)
class MCMC(Sampler): """ This class fits probability models using Markov Chain Monte Carlo. Each stochastic variable is assigned a StepMethod object, which makes it take a single MCMC step conditional on the rest of the model. These step methods are called in turn. >>> A = MCMC(input, db, verbose=0) :Parameters: - input : module, list, tuple, dictionary, set, object or nothing. Model definition, in terms of Stochastics, Deterministics, Potentials and Containers. If nothing, all nodes are collected from the base namespace. - db : string The name of the database backend that will store the values of the stochastics and deterministics sampled during the MCMC loop. - verbose : integer Level of output verbosity: 0=none, 1=low, 2=medium, 3=high Inherits all methods and attributes from Model. Subclasses must define the _loop method: - _loop(self, *args, **kwargs): Can be called after a sampling run is interrupted (by pausing, halting or a KeyboardInterrupt) to continue the sampling run. _loop must be able to handle KeyboardInterrupts gracefully, and should monitor the sampler's status periodically. Available status values are: - 'ready': Ready to sample. - 'paused': A pause has been requested, or the sampler is paused. _loop should return control as soon as it is safe to do so. - 'halt': A halt has been requested, or the sampler is stopped. _loop should call halt as soon as it is safe to do so. - 'running': Sampling is in progress. :SeeAlso: Model, Sampler, StepMethod. """ def __init__(self, input=None, db='ram', name='MCMC', calc_deviance=True, **kwds): """Initialize an MCMC instance. :Parameters: - input : module, list, tuple, dictionary, set, object or nothing. Model definition, in terms of Stochastics, Deterministics, Potentials and Containers. If nothing, all nodes are collected from the base namespace. - db : string The name of the database backend that will store the values of the stochastics and deterministics sampled during the MCMC loop. - verbose : integer Level of output verbosity: 0=none, 1=low, 2=medium, 3=high - **kwds : Keywords arguments to be passed to the database instantiation method. """ Sampler.__init__(self, input, db, name, calc_deviance=calc_deviance, **kwds) self._sm_assigned = False self.step_method_dict = {} for s in self.stochastics: self.step_method_dict[s] = [] self._state = [ 'status', '_current_iter', '_iter', '_tune_interval', '_burn', '_thin', '_tuned_count', '_tune_throughout', '_burn_till_tuned' ] def use_step_method(self, step_method_class, *args, **kwds): """ M.use_step_method(step_method_class, *args, **kwds) Example of usage: To handle stochastic A with a Metropolis instance, M.use_step_method(Metropolis, A, sig=.1) To subsequently get a reference to the new step method, S = M.step_method_dict[A][0] """ new_method = step_method_class(*args, **kwds) if self.verbose > 1: print_('Using step method %s. Stochastics: ' % step_method_class.__name__) for s in new_method.stochastics: self.step_method_dict[s].append(new_method) if self.verbose > 1: print_('\t' + s.__name__) if self._sm_assigned: self.step_methods.add(new_method) setattr(new_method, '_model', self) def remove_step_method(self, step_method): """ Removes a step method. """ try: for s in step_method.stochastics: self.step_method_dict[s].remove(step_method) if hasattr(self, "step_methods"): self.step_methods.discard(step_method) self._sm_assigned = False except AttributeError: for sm in step_method: self.remove_step_method(sm) def assign_step_methods(self, verbose=-1, draw_from_prior_when_possible=True): """ Make sure every stochastic variable has a step method. If not, assign a step method from the registry. """ if not self._sm_assigned: if draw_from_prior_when_possible: # Assign dataless stepper first last_gen = set([]) for s in self.stochastics - self.observed_stochastics: if s._random is not None: if len(s.extended_children) == 0: last_gen.add(s) dataless, dataless_gens = crawl_dataless( set(last_gen), [last_gen]) if len(dataless): new_method = DrawFromPrior(dataless, dataless_gens[::-1], verbose=verbose) setattr(new_method, '_model', self) for d in dataless: if not d.observed: self.step_method_dict[d].append(new_method) if self.verbose > 1: print_( 'Assigning step method %s to stochastic %s' % (new_method.__class__.__name__, d.__name__)) for s in self.stochastics: # If not handled by any step method, make it a new step method using the registry if len(self.step_method_dict[s]) == 0: new_method = assign_method(s, verbose=verbose) setattr(new_method, '_model', self) self.step_method_dict[s].append(new_method) if self.verbose > 1: print_('Assigning step method %s to stochastic %s' % (new_method.__class__.__name__, s.__name__)) self.step_methods = set() for s in self.stochastics: self.step_methods |= set(self.step_method_dict[s]) for sm in self.step_methods: if sm.tally: for name in sm._tuning_info: self._funs_to_tally[ sm._id + '_' + name] = lambda name=name, sm=sm: getattr(sm, name) else: # Change verbosity for step methods for sm_key in self.step_method_dict: for sm in self.step_method_dict[sm_key]: sm.verbose = verbose self.restore_sm_state() self._sm_assigned = True def sample(self, iter, burn=0, thin=1, tune_interval=1000, tune_throughout=True, save_interval=None, burn_till_tuned=False, stop_tuning_after=5, verbose=0, progress_bar=True): """ sample(iter, burn, thin, tune_interval, tune_throughout, save_interval, verbose, progress_bar) Initialize traces, run sampling loop, clean up afterward. Calls _loop. :Parameters: - iter : int Total number of iterations to do - burn : int Variables will not be tallied until this many iterations are complete, default 0 - thin : int Variables will be tallied at intervals of this many iterations, default 1 - tune_interval : int Step methods will be tuned at intervals of this many iterations, default 1000 - tune_throughout : boolean If true, tuning will continue after the burnin period (True); otherwise tuning will halt at the end of the burnin period. - save_interval : int or None If given, the model state will be saved at intervals of this many iterations - verbose : boolean - progress_bar : boolean Display progress bar while sampling. - burn_till_tuned: boolean If True the Sampler would burn samples until all step methods are tuned. A tuned step methods is one that was not tuned for the last `stop_tuning_after` tuning intervals. The burn-in phase will have a minimum of 'burn' iterations but could be longer if tuning is needed. After the phase is done the sampler will run for another (iter - burn) iterations, and will tally the samples according to the 'thin' argument. This means that the total number of iteration is update throughout the sampling procedure. If burn_till_tuned is True it also overrides the tune_thorughout argument, so no step method will be tuned when sample are being tallied. - stop_tuning_adfter: int the number of untuned successive tuning interval needed to be reach in order for the burn-in phase to be done (If burn_till_tuned is True). """ self.assign_step_methods(verbose=verbose) if burn > iter: raise ValueError( 'Burn interval cannot be larger than specified number of iterations.' ) self._n_tally = int(iter) - int(burn) if burn_till_tuned: self._stop_tuning_after = stop_tuning_after tune_throughout = False if verbose > 0: print "burn_til_tuned is True. tune_throughout is set to False" burn = int(max(burn, stop_tuning_after * tune_interval)) iter = self._n_tally + burn self._iter = int(iter) self._burn = int(burn) self._thin = int(thin) self._tune_interval = int(tune_interval) self._tune_throughout = tune_throughout self._burn_till_tuned = burn_till_tuned self._save_interval = save_interval length = max(int(np.floor((1.0 * iter - burn) / thin)), 1) self.max_trace_length = length # Flags for tuning self._tuning = True self._tuned_count = 0 # Progress bar self.pbar = None if not verbose and progress_bar: self.pbar = ProgressBar(self._iter) # Run sampler Sampler.sample(self, iter, length, verbose) def _loop(self): # Set status flag self.status = 'running' try: while self._current_iter < self._iter and not self.status == 'halt': if self.status == 'paused': break i = self._current_iter # Update progress bar if not (i + 1) % 100 and self.pbar: self.pbar.animate(i + 1) # Tune at interval if i and not (i % self._tune_interval) and self._tuning: self.tune() #update _burn and _iter if needed if self._burn_till_tuned and (self._tuned_count == 0): new_burn = self._current_iter + int( self._stop_tuning_after * self._tune_interval) self._burn = max(new_burn, self._burn) self._iter = self._burn + self._n_tally self.pbar = ProgressBar(self._iter) self.pbar.animate(i + 1) # Manage burn-in if i == self._burn: if self.verbose > 0: print_('\nBurn-in interval complete') if not self._tune_throughout: self._tuning = False # Tell all the step methods to take a step for step_method in self.step_methods: if self.verbose > 2: print_('Step method %s stepping' % step_method._id) # Step the step method step_method.step() # Record sample to trace, if appropriate if i % self._thin == 0 and i >= self._burn: self.tally() if self._save_interval is not None: if i % self._save_interval == 0: self.save_state() # Periodically commit samples to backend if not i % 1000: self.commit() # Increment interation self._current_iter += 1 except KeyboardInterrupt: self.status = 'halt' finally: # Stop progress bar if self.pbar: self.pbar.animate(self._iter) if self.status == 'halt': self._halt() def tune(self): """ Tell all step methods to tune themselves. """ # ======================================= # = This is what makes price.py puke... = # ======================================= # Only tune during burn-in # if self._current_iter > self._burn: # self._tuning = False # return if self.verbose > 0: print_('\tTuning at iteration', self._current_iter) # Initialize counter for number of tuning stochastics tuning_count = 0 for step_method in self.step_methods: verbose = self.verbose if step_method.verbose > -1: verbose = step_method.verbose # Tune step methods tuning_count += step_method.tune(verbose=self.verbose) if verbose > 1: print_('\t\tTuning step method %s, returned %i\n' % (step_method._id, tuning_count)) sys.stdout.flush() if self._burn_till_tuned: if not tuning_count: # If no step methods needed tuning, increment count self._tuned_count += 1 else: # Otherwise re-initialize count self._tuned_count = 0 # n consecutive clean intervals removed tuning # n is equal to self._stop_tuning_after if self._tuned_count == self._stop_tuning_after: if self.verbose > 0: print_('\nFinished tuning') self._tuning = False def get_state(self): """ Return the sampler and step methods current state in order to restart sampling at a later time. """ self.step_methods = set() for s in self.stochastics: self.step_methods |= set(self.step_method_dict[s]) state = Sampler.get_state(self) state['step_methods'] = {} # The state of each StepMethod. for sm in self.step_methods: state['step_methods'][sm._id] = sm.current_state().copy() return state def restore_sm_state(self): sm_state = self.db.getstate() if sm_state is not None: sm_state = sm_state.get('step_methods', {}) # Restore stepping methods state for sm in self.step_methods: sm.__dict__.update(sm_state.get(sm._id, {})) def _calc_dic(self): """Calculates deviance information Criterion""" # Find mean deviance mean_deviance = np.mean(self.db.trace('deviance')(), axis=0) # Set values of all parameters to their mean for stochastic in self.stochastics: # Calculate mean of paramter try: mean_value = np.mean(self.db.trace(stochastic.__name__)(), axis=0) # Set current value to mean stochastic.value = mean_value except KeyError: print_( "No trace available for %s. DIC value may not be valid." % stochastic.__name__) # Return twice deviance minus deviance at means return 2 * mean_deviance - self.deviance # Make dic a property def _get_dic(self): return self._calc_dic() dic = property(_get_dic)
def make_crystal(N_ions=6, starting_ions=None, constant_ion=None, progress=True, assymetry=None): ''' This function iterates random steps in postion and recalculates the energy of the configuration of N ions, if the energy is decreased it accepts the move and repeats ''' A = utility.get_potential_coefficient() tp = get_trap_parameters() sp = get_simulation_parameters() if assymetry: assy = assymetry else: assy = tp[4] pos_spread = sp[2] iterations = sp[1] step_size = sp[0] if progress: p = ProgressBar(iterations) N = N_ions ions = [] if starting_ions is None: for _i in range(N): # initiate ions in random position x0 = (np.random.rand() - 0.5)*pos_spread y0 = (np.random.rand() - 0.5)*pos_spread ions.append(ion(x0, y0)) else: # initiate ions in specified starting positions for starting_ion in starting_ions: ions.append(ion(starting_ion.x, starting_ion.y, starting_ion.color)) E0 = total_energy(ions, assy) iters = 0 E = [E0] last_E = E0 if constant_ion is not None: for fixed in constant_ion: ions[fixed].constant = True while iters < iterations: if progress: p.animate(iters) iters += 1 for particle in ions: if not particle.constant: # Checks if ion is movable particle.lastx = particle.x particle.x = particle.x + (np.random.rand() - 0.5)*step_size particle.lasty = particle.y particle.y = particle.y + (np.random.rand() - 0.5)*step_size tot_E = total_energy(ions, assy, A) if tot_E < last_E: # Check to take the position step or not E.append(tot_E) last_E = tot_E else: particle.x = particle.lastx particle.y = particle.lasty return E[-1], ions
def make_crystal(N_ions=6, starting_ions=None, constant_ion=None, progress=True, assymetry=None): ''' This function iterates random steps in postion and recalculates the energy of the configuration of N ions, if the energy is decreased it accepts the move and repeats ''' A = utility.get_potential_coefficient() tp = get_trap_parameters() sp = get_simulation_parameters() if assymetry: assy = assymetry else: assy = tp[4] pos_spread = sp[2] iterations = sp[1] step_size = sp[0] if progress: p = ProgressBar(iterations) N = N_ions ions = [] if starting_ions is None: for _i in range(N): # initiate ions in random position x0 = (np.random.rand() - 0.5) * pos_spread y0 = (np.random.rand() - 0.5) * pos_spread ions.append(ion(x0, y0)) else: # initiate ions in specified starting positions for starting_ion in starting_ions: ions.append(ion(starting_ion.x, starting_ion.y, starting_ion.color)) E0 = total_energy(ions, assy) iters = 0 E = [E0] last_E = E0 if constant_ion is not None: for fixed in constant_ion: ions[fixed].constant = True while iters < iterations: if progress: p.animate(iters) iters += 1 for particle in ions: if not particle.constant: # Checks if ion is movable particle.lastx = particle.x particle.x = particle.x + (np.random.rand() - 0.5) * step_size particle.lasty = particle.y particle.y = particle.y + (np.random.rand() - 0.5) * step_size tot_E = total_energy(ions, assy, A) if tot_E < last_E: # Check to take the position step or not E.append(tot_E) last_E = tot_E else: particle.x = particle.lastx particle.y = particle.lasty return E[-1], ions