Package: ratinabox
RatInABoxRatInABox(seepaper) is a toolkit for generating synthetic behaviour and neural data for spatially and/or velocity selective cell types in complex continuous environments.Install|Demos|Features|Contributions and Questions|CiteWithRatInABoxyou can:Generate realistic trajectoriesfor rats exploring complex 1 and 2D environments under a smooth random policy, an external control signal, or your own trajectory data.Generate artificial neuronal datafor various location- or velocity-selective cells found in the brain (e.g., but not limited to, Hippocampal cell types), or build your own more complex cell types.Build and train complex multi-layer networksof cells, powered by data generated withRatInABox.RatInABoxis an open source project welcomingcontributions. If you useRatInABoxpleasecitethe paper and consider giving this repository a star ☆. It contains three classes:Environment()📦: The environment/maze (or "box") that the agent lives in. 1- or 2-dimensional.Agent()🐀: The agent (or "rat") moving around theEnvironment.Neurons()🧠: A population of neurons with firing rates determined by the state (position and velocity) of theAgent. Make your own or use one of our premade cell types including:PlaceCells()GridCells()BoundaryVectorCells()(egocentric or allocentric)ObjectVectorCells()VelocityCells()SpeedCells()FieldOfViewNeurons()(egocentric encoding of what theAgentcan see)RandomSpatialNeurons()HeadDirectionCells()FeedForwardLayer()(a generic class analagous to a feedforward layer in a deep neural network)NeuralNetworkNeurons()(a generic class analagous to a deep neural network)SuccessorFeatures()...The top animation shows an example use case: anAgentrandomly explores a 2DEnvironmentwith a wall. Three populations ofNeurons(PlaceCells,GridCells,BoundaryVectorCells) fire according to the receptive fields shown. All data is saved into the history for downstream use.RatInABoxis fully continuous in space; this means that position and neural firing rates are calculated rapidly online with float precision rather than pre-calculated over a discretised mesh.RatInABoxis flexibly discretised in time;dtcan be set by the user (defaulting to 10 ms) depending on requirements.Key featuresNon-specific: Trajectories can be randomly generated, imported, or adaptively controlled makingRatInABoxa powerful engine for many tasks involving continuous motion (e.g. control theory orreinforcement learning).Biological:   Simulate large populations of spatially and/or velocity modulated cell types. Neurons can be rate based or spiking. The random motion model is fitted to match real rodent motion.Flexible:     Simulate environments in 1D or 2D with arbitrarily wall, boundary and hole arrangements.  Combine premade or bespokeNeuronsclasses into arbitrary deep networks (examples given).Fast:         Simulating 1 minute of exploration in a 2D environment with 100 place cells (dt=10 ms) take just 2 seconds on a laptop (no GPU needed).Precise:      No more prediscretised positions, tabular state spaces, or jerky movement policies. It's all continuous.Easy:         Sensible default parameters mean you can have realisitic simulation data to work with in <10 lines of code.VisualPlot or animate trajectories, firing rate timeseries', spike rasters, receptive fields, heat maps, velocity histograms...using the plotting functions (summarised here).Get startedManydemosare provided. Reading through theexample scripts(one simple and one extensive, duplicated at the bottom of the readme) these should be enough to get started. We also provide numerous interactive jupyter scripts as more in-depth case studies; for example one whereRatInABoxis used forreinforcement learning, another forneural decodingof position from firing rate. Jupyter scripts reproducing all figures in thepaperandreadmeare also provided. Alldemoscan be run on Google ColabInstalling and ImportingRequirementsare minimal (python3,numpy,scipyandmatplotlib, listed insetup.cfg) and will be installed automatically.Installthe latest, stable version usingpipat the command line with$pipinstallratinaboxAlternatively, in particular if you would like to developRatInABoxcode or if you want the bleeding edge (may occasioanlly break), install from this repo using$gitclone--depth1https://github.com/RatInABox-Lab/RatInABox.git$cdRatInABox$pipinstall-e.n.b. the "editable"-ehandle means changes made to your clone will be reflected when you next importRatInABoxinto your code.Importinto your python project withimportratinaboxfromratinabox.EnvironmentimportEnvironmentfromratinabox.AgentimportAgentfromratinabox.NeuronsimportPlaceCells,GridCells#...Feature run-downHere is a list of features loosely organised into those pertaining to(i) theEnvironmentAdding wallsComplex Environments: Polygons, curves and holesObjectsBoundary conditions1- or 2-dimensions(ii) theAgentRandom motionImporting trajectoriesPolicy controlWall repellingMultiple AgentsAdvancedAgentclasses(iii) theNeurons.Cell typesNoiseSpikes vs ratesPlotting rate mapsPlace cell modelsPlace cell geometryEgocentric encodingsReinforcement learning and successor featuresDeep neural networksCreate your ownNeurontypes(iv)Figures and animations plottingSpecific details can be found in thepaper.(i)EnvironmentfeaturesWallsArbitrarily add walls to the environment to produce arbitrarily complex mazes:Environment.add_wall([[x0,y0],[x1,y1]])Here are some easy to make examples.ComplexEnvironments: Polygons, curves, and holesBy default,Environmentsin RatInABox are square (or rectangular ifaspect != 1). It is possible to create arbitrary environment shapes using the"boundary"parameter at initialisation.You can also add holes to theEnvironmentusing the"holes"parameter at initialisation. When sampling positions from the Environment (e.g. at initialisation), holes won't be included.Any curved environments can be made by creating a boundary of many small walls (use sparingly, walls may slow down computations, particular for wall-responsive representations, e.g. boundary vector cells.)#A trapezium shaped EnvironmentEnv=Environment(params={'boundary':[[0,-0.2],[0,0.2],[1.5,0.5],[1.5,-0.5]],})#An environment with two holes making a figure of 8Env=Environment(params={'aspect':1.8,'holes':[[[0.2,0.2],[0.8,0.2],[0.8,0.8],[0.2,0.8]],[[1,0.2],[1.6,0.2],[1.6,0.8],[1,0.8]]],})#A circular environment made from many small wallsEnv=Environment(params={'boundary':[[0.5*np.cos(t),0.5*np.sin(t)]fortinnp.linspace(0,2*np.pi,100)],})ObjectsEnvironments can contain objects. These are used byObjectVectorCellsas visual cues but in theory you can hijack these to respresent many things in yourEnvironment(reward ports, goal locations etc.). Objects have atypeand apositionObjects are defined by a list of points and a position. Objects can be used as visual cues (e.g. forObjectVectorCells) or as obstacles (e.g. forBoundaryVectorCells).Envirnoment.add_object(object=[0.3,0.3],type=0)Envirnoment.add_object(object=[0.7,0.3],type=0)Envirnoment.add_object(object=[0.5,0.7],type=1)Boundary conditionsBoundary conditions (for default square/rectangular environments) can be "periodic" or "solid". Place cells and the motion of the Agent will respect these boundaries accordingly.Env=Environment(params={'boundary_conditions':'periodic'}#or 'solid' (default))1- or 2-dimensionsRatInABoxsupports 1- or 2-dimensionalEnvironments. Almost all applicable features and plotting functions work in both. The following figure shows 1 minute of exploration of anAgentin a 1D environment with periodic boundary conditions spanned by 10 place cells.Env=Environment(params={'dimensionality':'1D'}#or '2D' (default))(ii)AgentfeaturesRandom motion modelBy defaut theAgentfollows a random motion policy.  Random motion is stochastic but smooth. The speed (and rotational speed, if in 2D) of an Agent take constrained random walks governed by Ornstein-Uhlenbeck processes. You can change the means, variance and coherence times of these processes to control the shape of the trajectory. Default parameters are fit to real rat locomotion data from Sargolini et al. (2006):The default parameters can be changed to obtain different style trajectories. The following set of trajectories were generated by modifying the rotational speed parameterAgent.rotational_velocity_std:Agent.speed_mean=0.08#m/sAgent.speed_coherence_time=0.7Agent.rotation_velocity_std=120*np.pi/180#radiansAgent.rotational_velocity_coherence_time=0.08Importing trajectoriesRatInABoxsupports importing external trajectory data (rather than using the in built random motion policy). Imported data can be of low temporal resolution. It will be smoothly upsampled using a cubic splines interpolation technique. We provide a 10 minute trajectory from the open-source data set of Sargolini et al. (2006) ready to import. In the following figure blue shows (low resolution) trajectory data imported into anAgentand purple shows the smoothly upsampled trajectory taken by theAgentduring exploration.Agent.import_trajectory(dataset='sargolini')#orAgent.import_trajectory(times=array_of_times,positions=array_of_positions)Policy controlBy default the movement policy is an random and uncontrolled (e.g. displayed above). It is possible, however, to manually pass a "drift_velocity" to the Agent on eachAgent.update()step. This 'closes the loop' allowing, for example, Actor-Critic systems to control the Agent policy. As a demonstration that this method can be used to control the agent's movement we set a radial drift velocity to encourage circular motion. We also use RatInABox to perform a simple model-free RL task and find a reward hidden behind a wall (the full script is given as an example scripthere)Agent.update(drift_velocity=drift_velocity)Wall repellingUnder the random motion policy, walls in the environment mildly "repel" theAgent. Coupled with the finite turning speed this replicates an effect (known as thigmotaxis, sometimes linked to anxiety) where theAgentis biased to over-explore near walls and corners (as shown in these heatmaps) matching real rodent behaviour. It can be turned up or down with thethigmotaxisparameter.Αgent.thigmotaxis=0.8#1 = high thigmotaxis (left plot), 0 = low (right)MultipleAgentsThere is nothing to stop multipleAgents being added to the sameEnvironment. When plotting/animating trajectories set the kwargplot_all_agents=Trueto visualise allAgents simultaneously.The following animation shows threeAgents in anEnvironment. Drift velocities are set so thatAgents weally locally attract one another creating "interactive" behaviour.AdvancedAgentclassesOne can make more advanced Agent classes, for exampleThetaSequenceAgent()where the position "sweeps" (blue) over the position of an underlying true (regular)Agent()(purple), highly reminiscent of theta sequences observed when one decodes position from the hippocampal populaton code on sub-theta (10 Hz) timescales. This class can be found in thecontribsdirectory.(iii)NeuronsfeaturesMultiple cell types:We provide a list of premadeNeuronssubclasses. These include (but are not limited to):PlaceCellsGridCellsBoundaryVectorCells(can be egocentric or allocentric)ObjectVectorCells(can be used as visual cues, i.e. only fire whenAgentis looking towards them) (can be egocentric or allocentric)HeadDirectionCellsVelocityCellsSpeedCellsRandomSpatialNeurons- smooth but random spatially tuned neuronsFeedForwardLayer- calculates activated weighted sum of inputs from a provide list of inputNeuronslayers.FieldOfViewNeurons- Egocentric encoding of what theAgentcan seeNeuralNetworkNeurons- Maps inputs from a user provided list of inputNeuronsthrough a user-providedpytorchneural network. Can be used to create arbitrary and learnable representations.FeedForwardLayerandNeuralNetworkNeuronsdeserves special mention. Instead of its firing rate being determined explicitly by the state of theAgentit summates synaptic inputs from a provided list of input layers (which can be anyNeuronssubclass). This layer is the building block for how more complex networks can be studied usingRatInABox.NeuralNetworkNeuronsis the same except instead of linearly summating it passes inputs through any arbitrary deep neural network.NoiseUse theNeurons.noise_stdandNeurons.noise_coherence_timeparameters to control the amount of noise (Hz) and autocorrelation timescale of the noise (seconds). For example (work with allNeuronsclasses, not justPlaceCells):PCs=PlaceCells(Ag,params={'noise_std':0.1,#defaults to 0 i.e. no noise'noise_coherence_time':0.5,#autocorrelation timescale of additive noise vector})SpikingAll neurons are rate based. However, at each update spikes are sampled as though neurons were Poisson neurons. These are stored inNeurons.history['spikes']. The max and min firing rates can be set withNeurons.max_frandNeurons.min_fr.Neurons.plot_ratemap(spikes=True)Rate mapsPlaceCells,GridCellsand allocentricBoundaryVectorCells(among others) have firing rates which depend exclusively on the position of the agent. These rate maps can be displayed by querying their firing rate at an array of positions spanning the environment, then plotting. This process is done for you using the functionNeurons.plot_rate_map().More generally, however, cells firing is not only determined by position but potentially other factors (e.g. velocity, or historical effects if the layer is part of a recurrent network). In these cases the above method of plotting rate maps will necessarily fail. A more robust way to display the receptive field is to plot a heatmap of the positions of the Agent has visited where each positions contribution to a bin is weighted by the firing rate observed at that position. Over time, as coverage become complete, the firing fields become visible.Neurons.plot_rate_map() #attempts to plot "ground truth" rate map 
Neurons.plot_rate_map(method="history") #plots rate map by firing-rate-weighted position heatmapPlace cell modelsPlace cells come in multiple types (given byparams['description']), or it would be easy to write your own:"gaussian": normal gaussian place cell"gaussian_threshold": gaussian thresholded at 1 sigma"diff_of_gaussians": gaussian(sigma) - gaussian(1.5 sigma)"top_hat": circular receptive field, max firing rate within, min firing rate otherwise"one_hot": the closest place cell to any given location is established. This and only this cell fires.This last place cell type,"one_hot"is particularly useful as it essentially rediscretises space and tabularises the state space (gridworld again). This can be used to  contrast and compare learning algorithms acting over continuous vs discrete state spaces. This figure compares the 5 place cell models for population of 9 place cells (top left shows centres of place cells, and in all cases the"widths"parameters is set to  0.2 m, or irrelevant in the case of"one_hot"s)These place cells (with the exception of"one_hot"s) can all be made to phase precess by instead initialising them with thePhasePrecessingPlaceCells()class currently residing in thecontribsfolder. This figure shows example output data.Geometry ofPlaceCellsChoose how you wantPlaceCellsto interact with walls in theEnvironment. We provide three types of geometries.Egocentric encodingsMostRatInABoxcell classes are allocentric (e.g.PlaceCells,GridCellsetc. do not depend on the agents point of view) not egocentric.BoundaryVectorCells(BVCs) andObjectVectorCells(OVCs) can be either.FieldOfViewNeuronsexploit this by arranging sets of egocentric BVC or OVCs to tile to agents local field of view creating a comprehensive egocentric encoding of what boundaries or objects the agent can 'see' from it's current point of view. A custom plotting function displays the tiling and the firing rates as shown below. With an adequately defined field of view these can make, for example, "whisker cells".FoV_BVCs=FieldOfViewBVCs(Ag)FoV_OVCs=FieldOfViewOVCs(Ag)BVCs_whiskers=FieldOfViewBVCs(Ag,params={"distance_range":[0.01,0.2],"angle_range":[75,105,],"spatial_resolution":0.02,# resolution of each OVC tiling FoV"cell_arrangement":"uniform_manifold",})Reinforcement Learning and Successor FeaturesA dedicatedNeuronsclass calledSuccessorFeatureslearns the successor features for a given feature set under the current policy. Seethis demofor more info.SuccessorFeaturesare a specific instance of a more general class of neurons calledValueNeurons which learn value function for any reward density under theAgents motion policy. This can be used to do reinforcement learning tasks such as finding rewards hidden behind walls etc as shown inthis demo.We also have a working examples of an actor critic algorithm using deep neural networkshereFinally, we are working on a dedicated subpackage -- (RATS: RL Agent Toolkit and Simulator) -- to host all this RL stuff and more so keep  an eye out.Neuronsas function approximatorsPerhaps you want to generate really complex cell types (more complex than justPlaceCells,GridCellsetc.). No problem. For this we provide twoNeuronssubclass useful for constructing complex Neuron types. For these classes instead of firing rates being determined explicitly by the state of theAgent, they recieve inputs from one or many otherRatInABox.Neuronsclasses and pass these inputs through a function to calculate the firing rate.FeedForwardLayerlinearly sums their inputs with a set of weights.NeuralNetworkNeuronsare more general, they pass their inputs through a user-provided deep neural network (for this we usepytorch).Both of these classes can be used as the building block for constructing complex multilayer networks ofNeurons(e.g.FeedForwardLayersareRatInABox.Neuronsin their own right so can be used as inputs to otherFeedForwardLayers). Their parameters can be accessed and set (or "trained") to create neurons with complex receptive fields. In the case ofDeepNeuralNetworkneurons the firing rate attached to the computational graph is stored so gradients can be taken. Examples can be foundhere (deep learning withNeuralNewtworkNeurons)as well ashere (path integration),here (reinforcement learning),here (successor features)andhere (actor critic deep RL).Creating your ownNeurontypesWe encourage users to create their own subclasses ofNeurons. This is easy to do, see comments in theNeuronsclass within thecodefor explanation. By forming these classes from the parentNeuronsclass, the plotting and analysis features described above remain available to these bespoke Neuron types.(iv) Figures and animationsRatInABoxis built to be highly visual. It is easy to plot or animate data and save these plots/animations. Here are some tipsSavingratinabox.figure_directorya global variable specifying the directory into which figures/animations will be savedratinabox.utils.save_figure(fig,fig_name)saves a figure (or animation) into a dated folder within the figure directory  as both".svg"and".png"(".mp4"or".gif"). The current time will be appended to thefig_nameso you won't ever overwrite.Saving (but automatically)Settingratinabox.autosave_plots = Truemeans RatInABox figure will be automatically saved in the figure directory without having to indvidually call theutilsfunction above.Stylingratinabox.stylize_plots()this call sets some global matplotlib rcParams to make plots look pretty/exactly like they do in this repoMost important plotting functionsThe most important plotting functions are (see source code for the available arguments/kwargs):Environment.plot_environment()#visualises current environment with walls and objectsAgent.plot_trajectory()#plots trajectoryAgent.animate_trajectory()#animate trajectoryNeurons.plot_rate_map()# plots the rate map of the neurons at all positionsNeurons.plot_rate_timeseries()# plots activities of the neurons over timeNeurons.animate_rate_timeseries()# animates the activity of the neurons over timeMost plotting functions acceptfigandaxas optional arguments and if passed will plot ontop of these. This can be used to make comolex or multipanel figures. For a comprehensive list of plotting functions seehere.Example ScriptsIn the folder calleddemoswe provide numerous script and demos which will help when learningRatInABox. In approximate order of complexity, these include:simple_example.ipynb: a very simple tutorial for importing RiaB, initialising an Environment, Agent and some PlaceCells, running a brief simulation and outputting some data. Code copied here for convenience.importratinabox#IMPORTfromratinabox.EnvironmentimportEnvironmentfromratinabox.AgentimportAgentfromratinabox.Neuronsimport*#INITIALISE CLASSESEnv=Environment()Ag=Agent(Env)PCs=PlaceCells(Ag)#EXPLOREforiinrange(int(20/Ag.dt)):Ag.update()PCs.update()#ANALYSE/PLOTprint(Ag.history['pos'][:10])print(PCs.history['firingrate'][:10])fig,ax=Ag.plot_trajectory()fig,ax=PCs.plot_rate_timeseries()extensive_example.ipynb: a more involved tutorial. More complex enivornment, more complex cell types and more complex plots are used.list_of_plotting_functions.md: All the types of plots available for are listed and explained.readme_figures.ipynb: (Almost) all plots/animations shown in the root readme are produced from this script (plus some minor formatting done afterwards in powerpoint).paper_figures.ipynb: (Almost) all plots/animations shown in the paper are produced from this script (plus some major formatting done afterwards in powerpoint).decoding_position_example.ipynb: Postion is decoded from neural data generated with RatInABox. Place cells, grid cell and boundary vector cells are compared.conjunctive_gridcells_example.ipynb:GridCellsandHeadDirectionCellsare minimally combined useingFeedForwardLayerto create head-direction-selective grid cells (aka. conjunctive cells).splitter_cells_example.ipynb: A simple simultaion demonstrating howSplittrecell data could be create in a figure-8 maze.deep_learning_example.ipynb: Here we showcaseNeuralNetworkNeuronsa class ofNeuronswhich has a small neural network embedded inside. We train them to take grid cells as inputs and output an arbitrary function  as their rate map.reinforcement_learning_example.ipynb: RatInABox is use to construct, train and visualise a small two-layer network capable of model free reinforcement learning in order to find a reward hidden behind a wall.actor_critic_example.ipynb: RatInABox is use to implement the actor critic algorithm using deep neural networks.successor_features_example.ipynb: RatInABox is use to learn and visualise successor features under random and biased motion policies.path_integration_example.ipynb: RatInABox is use to construct, train and visualise a large multi-layer network capable of learning a "ring attractor" capable of path integrating a position estimate using only velocity inputs.ContributeRatInABoxis open source project and we actively encourage  all contributions from example bug fixes to documentation or new cell types. Feel free to make a pull request (you will need to fork the repository first) or raise and issue.We have a dedicatedcontribsdirectory where you can safely add awesome scripts and newNeuronsclasses etc.Questions? Just ask! Ideally via opening an issue so others can see the answer too.Thanks to all contributors so far:CiteIf you useRatInABoxin your research or educational material, please cite the work as follows:
Bibtex:@article{ratinabox2022,
	doi = {10.1101/2022.08.10.503541},
	url = {https://doi.org/10.1101%2F2022.08.10.503541},
	year = 2022,
	month = {aug},
	publisher = {Cold Spring Harbor Laboratory},
	author = {Tom M George and William de Cothi and Claudia Clopath and Kimberly Stachenfeld and Caswell Barry},
	title = {{RatInABox}: A toolkit for modelling locomotion and neuronal activity in continuous environments}
}Formatted:Tom M George, William de Cothi, Claudia Clopath, Kimberly Stachenfeld, Caswell Barry. "RatInABox: A toolkit for modelling locomotion and neuronal activity in continuous environments" (2022).The research paper corresponding to the above citation can be foundhere.