RatInABox-Lab / RatInABox

A python package for modelling locomotion in complex environments and spatially/velocity selective cell activity.
MIT License
168 stars 31 forks source link

straight trajectory with fixed speed mean, but Ag.history['vel'] is not constant #113

Open ZilongJi opened 2 weeks ago

ZilongJi commented 2 weeks ago

Hi there, I have a question about the velocity extracted from the Agent. Basically, I simply set speed_mean=1, speed_std=0, and rotational_velocity_std = 0. I expected the trajectory is straight which is ture after I did the visualziation. However, I also expect the speed to be constant, but it is not the case. Please let me know if I should also set other parameters to get the correct speed at each moment. Many thanks!

image

Below is my code:


from ratinabox.Environment import Environment from ratinabox.Agent import Agent import numpy as np import matplotlib.pyplot as plt Env = Environment(params={ 'scale':2*np.pi, "dimensionality": "2D", })

dur = 10 #s dt = 0.001 #ms speed_mean = 1 #in m/ms speed_std = 0.0#in m/ms rotational_velocity_std = 0.0* (np.pi / 180) #in rad/s

Ag = Agent(Env, params = { "dt": dt, "speed_mean":speed_mean, "speed_std": speed_std, "rotational_velocity_std": rotational_velocity_std, #set to 120 or 360 will change the tutoriocity of the trajectory })

while Ag.t < dur: Ag.update()

Position = Ag.history['pos'] HeadDirection = Ag.history['head_direction'] Velocity = Ag.history['vel']

Position = np.array(Position) HeadDirection = np.array(HeadDirection) Velocity = np.array(Velocity)

plot trajectory

fig, axs = plt.subplots(1,3,figsize=(6,2)) ax = axs[0] ax.plot(Position[:,0], Position[:,1]) ax.axis('equal')
ax.set_title('Trajectory')

plot velocity

ax = axs[1] ax.plot(np.linalg.norm(Velocity, axis=1)) ax.set_title('Speed')

plot HeadDirection

ax = axs[2] HeadDirection = np.array(HeadDirection) complex_HD = HeadDirection[:, 0] + 1j * HeadDirection[:, 1] HD_angle = np.angle(complex_HD) ax.plot(HD_angle) ax.set_title('Head Direction')

plt.tight_layout()

TomGeorge1234 commented 1 week ago

Hi Zilong,

Why does this happen

This behaviour occurs because, in 2D, speed is sampled from a Rayleigh distribution where the scale parameter == speed_mean. The scale parameter for Rayleigh is, in a sense, the mean and std. It's a bit like a continuous poisson.

In fact, speed_std is totally ignored in 2D. This is a slighty confusing API but was chosen because the Rayleigh distribution is positive definite so speeds naturally don't go negative (also fits natural foraging really well). You'll find it detailed in the paper and/or docs:

https://github.com/RatInABox-Lab/RatInABox/blob/e8ec06294432ed1c82d3ae277a9ccc4876a8c263/ratinabox/Agent.py#L76

So how to get around this?

Surprisingly I've never actually thought about your fairly simple and reasonable requirement so a fix isn't completely natural. I can think of three way on the quick and hacky to formalised spectrum:

Easy: You could just overwrite the velocity on each update step

while Ag.t < dur:
    Ag.velocity = speed_mean * Ag.velocity / np.linalg.norm(Ag.velocity)
    Ag.update()
Screenshot 2024-06-24 at 11 17 16

You'll still see small variation because the stochastic update is still performed on each step, they just won't accumulate over time.

Medium: If you don't even want that small variance, get past this you could actually edit the source code (if you're building off a fork) and add the following two lines of code

if self.speed_std == 0:
    speed_new = speed_mean

between lines 309 and 310 of the Agent.py source code. https://github.com/RatInABox-Lab/RatInABox/blob/e8ec06294432ed1c82d3ae277a9ccc4876a8c263/ratinabox/Agent.py#L298-L310

image

which essentially overwrites the stochastic update. It's a bit hacky and I'm not sure yet if I'd want this added to the source code.

Hard The best long term solution would be to consider updating the motion model to replace Rayleigh with a different distribution which uses speed_std or formalise the above two lines of code into the API and document it clearly. Happy to consider either if you feel strongly but I reckon medium is your best bet...

TomGeorge1234 commented 1 week ago

oh and one last thing, turn wall_repel_strength to zero and agents will just bounce off walls so you'll also lose the variance near bounces.

Ag = Agent(Env, params = {
"dt": dt,
"speed_mean":speed_mean,
"speed_std": speed_std,
"rotational_velocity_std": rotational_velocity_std, #set to 120 or 360 will change the tutoriocity of the trajectory
"wall_repel_strength":0.0,
})
image
TomGeorge1234 commented 1 week ago

I've thought a little more about this. If I ever get round to it it might be nice to update the speed distribution to a Gamma distribution (similar to Rayleigh but includes a variance parameter.) I have no plans to do this soon so the above fixes are still your best option but I'm noting it here for future me.

$$\textrm{speed} \sim \Gamma(k,\theta)$$

The transforms would be:

$$k = \frac{\mu^2}{\sigma^2}$$

$$\theta = \frac{\sigma^2}{\mu}$$

though I should check how stable this is in the limit that $\sigma$ gets small.

TomGeorge1234 commented 1 week ago

@ZilongJi has this been resolved?