How to treat behavioral variables for CEBRA behavior

[GioLetti]

Hi,

thanks for the super interesting and innovative paper. I have a quick question regarding the usage of CEBRA behavior. I'm using CEBRA to perform manifold learning on calcium imaging data during freely moving navigation, but I don't know exactly how CEBRA handles certain behavioral variables or how I should treat them.

Specifically, I have the x-y mouse position on a 2D arena and the allocentric head direction in radians. Can I pass to the fit function as labels the data as a matrix [time_points, n_dims] with n_dims being x pos, y_pos, and head direction in radians? Or should I bin the variables and provide a linearized vector of positional bin number and head direction bin number?

Thanks for your amazing work,

Giorgio

[stes]

Hi @GioLetti , as a quick first start I would recommend to use x, y, cos direction, sin direction. The reason is that the sampling assumes an approximately Gaussian distribution around the passed variables by default.

The more sophisticated option would be to implement a custom sampler that can handle angles. But I think for a first impression, the suggestion above is an acceptable pragmatic solution.

[stes]

Please see here

CEBRA/cebra/distributions/continuous.py

Lines 254 to 298 in d1842cc

	class DeltaNormalDistribution(abc_.JointDistribution, abc_.HasGenerator):
	"""Define a conditional distribution based on behavioral changes over time.
	Takes a continuous index, and uses sample from Gaussian distribution to sample positive pairs.
	Note that if the continuous index is multidimensional, the Gaussian distribution will have
	isotropic covariance matrix i.e. Σ = sigma^2 * I.
	Args:
	continuous: The multidimensional, continuous index.
	delta: Standard deviation of Gaussian distribution to sample positive pair.
	"""
	def __init__(self,
	continuous: torch.Tensor,
	delta: float = 0.1,
	device: Literal["cpu", "cuda"] = "cpu",
	seed: Optional[int] = None):
	abc_.HasGenerator.__init__(self, device=device, seed=seed)
	self.data = continuous
	self.std = delta
	self.index = cebra.distributions.ContinuousIndex(self.data)
	self.prior = Prior(self.data, device=device, seed=seed)
	def sample_prior(self, num_samples: int) -> torch.Tensor:
	"""See :py:meth:`.Prior.sample_prior`."""
	return self.prior.sample_prior(num_samples)
	def sample_conditional(self, reference_idx: torch.Tensor) -> torch.Tensor:
	"""Return indices from the conditional distribution."""
	if reference_idx.dim() != 1:
	raise ValueError(
	f"Reference indices have wrong shape: {reference_idx.shape}. "
	"Pass a 1D array of indices of reference samples.")
	# TODO(stes): Set seed
	mean = self.data[reference_idx]
	query = torch.distributions.Normal(
	loc=mean,
	scale=torch.ones_like(mean, device=self.device) * self.std,
	).sample()
	query = query.unsqueeze(-1) if query.dim() == 1 else query
	return self.index.search(query)

and here

CEBRA/cebra/distributions/continuous.py

Lines 200 to 238 in d1842cc

	class TimedeltaDistribution(abc_.JointDistribution, abc_.HasGenerator):
	"""Define a conditional distribution based on behavioral changes over time.
	Takes a continuous index, and uses the empirical distribution of differences
	between samples in this index.
	Args:
	continuous: The multidimensional, continuous index
	time_delta: The time delay between samples that should form a positive
	pair.
	device: TODO
	seed: TODO
	Note:
	For best results, the given continuous index should contain independent
	factors; positive pairs will be formed by adding a _random_ difference
	estimated within the dataset to the reference samples.
	Factors should ideally also be within the same range (since the Euclidean
	distance is used in the search). A simple solution is to perform a PCA
	or ICA, or apply CEBRA first before passing the index to this function.
	"""
	def __init__(self,
	continuous,
	time_delta: int = 1,
	device: Literal["cpu", "cuda"] = "cpu",
	seed: Optional[int] = None):
	abc_.HasGenerator.__init__(self, device=device, seed=seed)
	self.data = continuous
	self.time_delta = time_delta
	self.time_difference = torch.zeros_like(self.data, device=self.device)
	self.time_difference[time_delta:] = (self.data[time_delta:] -
	self.data[:-time_delta])
	self.index = cebra.distributions.ContinuousIndex(self.data)
	self.prior = Prior(self.data, device=device, seed=seed)
	def sample_prior(self, num_samples: int) -> torch.Tensor:
	"""See :py:meth:`.Prior.sample_prior`."""
	return self.prior.sample_prior(num_samples)

[stes]

It could be an interesting contribution to the package to add a distribution that properly handles x,y and the head direction.