kalmax 0.0.0__tar.gz

kalmax-0.0.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2022 Tom M George
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

kalmax-0.0.0/PKG-INFO

@@ -0,0 +1,123 @@
+Metadata-Version: 2.1
+Name: kalmax
+Version: 0.0.0
+Summary: Kalman based neural decoding in Jax
+Home-page: https://github.com/TomGeorge1234/KalMax
+Author: Tom George
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.6
+Description-Content-Type: text/markdown
+Provides-Extra: demo
+License-File: LICENSE
+
+# **KalMax**:  Kalman based neural decoding in Jax 
+**KalMax** = **Kal**man smoothing of **Max**imum likelihood estimates in Jax.
+
+You provide $\mathbf{S} \in \mathbb{N}^{T \times N}$ (spike counts) and $\mathbf{X} \in \mathbb{R}^{T \times D}$ (a continuous variable, e.g. position) and `KalMax` provides jax-optimised functions and classes for:
+
+1. **Fitting rate maps** using kernel density estimation (KDE)
+2. **Calculating likelihood** maps $p(\mathbf{s}_t|\mathbf{x})$
+3. **Kalman filter / smoother**
+
+<img src="figures/display_figures/input_data.png" width=350>
+
+
+
+
+#### Why are these functionalities combined into one package?...
+
+Because Likelihood Estimation + Kalman filtering = Powerful neural decoding. By Kalman filtering/smoothing the maximum likelihood estimates (as opposed to the spikes themselves) we bypass the issues of naive Kalman filters (spikes are rarely linearly related to position) and maximum likelihood decoding (which does not account for temporal continuity in the trajectory), outperforming both for no extra computational cost.
+<img src="figures/display_figures/filter_comparisons.gif" width=850>
+
+
+Core `KalMax` functions are optimised and jit-compiled in jax making them **very fast**. For example `KalMax` kalman filtering is >13 times faster than an equivalent numpy implementation by the popular [`pykalman`](https://github.com/pykalman/pykalman/tree/master) library (see [demo](./kalmax_demo.ipynb)).
+
+<img src="figures/display_figures/kalman_speed_comparison.png" width=150>
+
+
+# Install
+```
+git clone https://github.com/TomGeorge1234/KalMax.git
+cd KalMax
+pip install -e .
+```
+(`-e`) is optional for developer install. 
+
+Alternatively 
+```
+pip install git+https://github.com/TomGeorge1234/KalMax.git
+```
+
+# Usage  
+
+A full demo [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/TomGeorge1234/KalMax/blob/main/kalmax_demo.ipynb) is provided in the [`kalmax_demo.ipynb`](./kalmax_demo.ipynb). Sudo-code is provided below. 
+
+```python
+import kalmax 
+import jax.numpy as jnp 
+```
+
+```python
+# 0. PREPARE DATA IN JAX ARRAYS
+S_train = jnp.array(...) # (T, N_CELLS)      train spike counts
+Z_train = jnp.array(...) # (T, DIMS)         train continuous variable
+S_test  = jnp.array(...) # (T_TEST, N_CELLS) test spike counts
+bins    = jnp.array(...) # (N_BINS, DIMS)    coordinates at which to estimate receptive fields / likelihoods)
+```
+<img src="figures/display_figures/data.png" width=850>
+
+```python
+# 1. FIT RECEPTIVE FIELDS using kalmax.kde
+firing_rate = kalmax.kde.kde(
+    bins = bins,
+    trajectory = Z_train,
+    spikes = S_train,
+    kernel = kalmax.kernels.gaussian_kernel,
+    kernel_kwargs = {'covariance':0.01**2*np.eye(DIMS)}, # kernel bandwidth
+    ) # --> (N_CELLS, N_BINS)
+```
+<img src="figures/display_figures/receptive_fields.png" width=850>
+
+
+```python
+# 2.1 CALCULATE LIKELIHOODS using kalmax.poisson_log_likelihood
+log_likelihoods = kalmax.kde.poisson_log_likelihood(
+    spikes = S_test,                       
+    mean_rate = firing_rate,
+    ) # --> (T_TEST, N_CELLS)
+
+# 2.2 FIT GAUSSIAN TO LIKELIHOODS using kalmax.utils.fit_gaussian
+MLE_means, MLE_modes, MLE_covs = kalmax.utils.fit_gaussian_vmap(
+    x = bins, 
+    likelihoods = jnp.exp(log_likelihoods),
+    ) # --> (T_TEST, DIMS), (T_TEST, DIMS, DIMS)
+```
+<img src="figures/display_figures/likelihood_maps_fitted.png" width=850>
+
+```python
+# 3. KALMAN FILTER / SMOOTH using kalmax.KalmanFilter.KalmanFilter
+kalman_filter = kalmax.kalman.KalmanFilter(
+    dim_Z = DIMS, 
+    dim_Y = N_CELLS,
+    # SEE DEMO FOR HOW TO FIT/SET THESE
+    F=F, # state transition matrix
+    Q=Q, # state noise covariance
+    H=H, # observation matrix
+    R=R, # observation noise covariance
+    ) 
+
+# [FILTER]
+mus_f, sigmas_f = kalman_filter.filter(
+    Y = Y, 
+    mu0 = mu0,
+    sigma0 = sigma0,
+    ) # --> (T, DIMS), (T, DIMS, DIMS)
+
+# [SMOOTH]
+mus_s, sigmas_s = kalman_filter.smooth(
+    mus_f = mus_f, 
+    sigmas_f = sigmas_f,
+    ) # --> (T, DIMS), (T, DIMS, DIMS)
+```

kalmax-0.0.0/README.md

@@ -0,0 +1,109 @@
+# **KalMax**:  Kalman based neural decoding in Jax 
+**KalMax** = **Kal**man smoothing of **Max**imum likelihood estimates in Jax.
+
+You provide $\mathbf{S} \in \mathbb{N}^{T \times N}$ (spike counts) and $\mathbf{X} \in \mathbb{R}^{T \times D}$ (a continuous variable, e.g. position) and `KalMax` provides jax-optimised functions and classes for:
+
+1. **Fitting rate maps** using kernel density estimation (KDE)
+2. **Calculating likelihood** maps $p(\mathbf{s}_t|\mathbf{x})$
+3. **Kalman filter / smoother**
+
+<img src="figures/display_figures/input_data.png" width=350>
+
+
+
+
+#### Why are these functionalities combined into one package?...
+
+Because Likelihood Estimation + Kalman filtering = Powerful neural decoding. By Kalman filtering/smoothing the maximum likelihood estimates (as opposed to the spikes themselves) we bypass the issues of naive Kalman filters (spikes are rarely linearly related to position) and maximum likelihood decoding (which does not account for temporal continuity in the trajectory), outperforming both for no extra computational cost.
+<img src="figures/display_figures/filter_comparisons.gif" width=850>
+
+
+Core `KalMax` functions are optimised and jit-compiled in jax making them **very fast**. For example `KalMax` kalman filtering is >13 times faster than an equivalent numpy implementation by the popular [`pykalman`](https://github.com/pykalman/pykalman/tree/master) library (see [demo](./kalmax_demo.ipynb)).
+
+<img src="figures/display_figures/kalman_speed_comparison.png" width=150>
+
+
+# Install
+```
+git clone https://github.com/TomGeorge1234/KalMax.git
+cd KalMax
+pip install -e .
+```
+(`-e`) is optional for developer install. 
+
+Alternatively 
+```
+pip install git+https://github.com/TomGeorge1234/KalMax.git
+```
+
+# Usage  
+
+A full demo [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/TomGeorge1234/KalMax/blob/main/kalmax_demo.ipynb) is provided in the [`kalmax_demo.ipynb`](./kalmax_demo.ipynb). Sudo-code is provided below. 
+
+```python
+import kalmax 
+import jax.numpy as jnp 
+```
+
+```python
+# 0. PREPARE DATA IN JAX ARRAYS
+S_train = jnp.array(...) # (T, N_CELLS)      train spike counts
+Z_train = jnp.array(...) # (T, DIMS)         train continuous variable
+S_test  = jnp.array(...) # (T_TEST, N_CELLS) test spike counts
+bins    = jnp.array(...) # (N_BINS, DIMS)    coordinates at which to estimate receptive fields / likelihoods)
+```
+<img src="figures/display_figures/data.png" width=850>
+
+```python
+# 1. FIT RECEPTIVE FIELDS using kalmax.kde
+firing_rate = kalmax.kde.kde(
+    bins = bins,
+    trajectory = Z_train,
+    spikes = S_train,
+    kernel = kalmax.kernels.gaussian_kernel,
+    kernel_kwargs = {'covariance':0.01**2*np.eye(DIMS)}, # kernel bandwidth
+    ) # --> (N_CELLS, N_BINS)
+```
+<img src="figures/display_figures/receptive_fields.png" width=850>
+
+
+```python
+# 2.1 CALCULATE LIKELIHOODS using kalmax.poisson_log_likelihood
+log_likelihoods = kalmax.kde.poisson_log_likelihood(
+    spikes = S_test,                       
+    mean_rate = firing_rate,
+    ) # --> (T_TEST, N_CELLS)
+
+# 2.2 FIT GAUSSIAN TO LIKELIHOODS using kalmax.utils.fit_gaussian
+MLE_means, MLE_modes, MLE_covs = kalmax.utils.fit_gaussian_vmap(
+    x = bins, 
+    likelihoods = jnp.exp(log_likelihoods),
+    ) # --> (T_TEST, DIMS), (T_TEST, DIMS, DIMS)
+```
+<img src="figures/display_figures/likelihood_maps_fitted.png" width=850>
+
+```python
+# 3. KALMAN FILTER / SMOOTH using kalmax.KalmanFilter.KalmanFilter
+kalman_filter = kalmax.kalman.KalmanFilter(
+    dim_Z = DIMS, 
+    dim_Y = N_CELLS,
+    # SEE DEMO FOR HOW TO FIT/SET THESE
+    F=F, # state transition matrix
+    Q=Q, # state noise covariance
+    H=H, # observation matrix
+    R=R, # observation noise covariance
+    ) 
+
+# [FILTER]
+mus_f, sigmas_f = kalman_filter.filter(
+    Y = Y, 
+    mu0 = mu0,
+    sigma0 = sigma0,
+    ) # --> (T, DIMS), (T, DIMS, DIMS)
+
+# [SMOOTH]
+mus_s, sigmas_s = kalman_filter.smooth(
+    mus_f = mus_f, 
+    sigmas_f = sigmas_f,
+    ) # --> (T, DIMS), (T, DIMS, DIMS)
+```

kalmax-0.0.0/kalmax/demo_utils.py

@@ -0,0 +1,119 @@
+# Some (mostly plotting and poorly written) utilities for the Kalman filter demo
+import matplotlib
+import matplotlib.pyplot as plt
+import numpy as np
+
+def animate_over_time(
+            plot_function,
+            t_start,
+            t_end,
+            speed_up=10, #relative to real time
+            fps=20,):
+    """Animates the plot function over time_range.
+
+    Params
+    ------
+    plot_function: function, it must act as it's own init function when ax=None
+        Function that takes ax, time and time_start as input and returns ax.
+        Must have the following signature 
+        def plot_function(ax,
+                            time,
+                            time_start=0 # time to start (relevant if plotting a trajectory) 
+                            ):
+            # do something
+            return ax
+        _type_: _description_
+    """
+    # plot function should take ax and time as input
+    time_per_frame = speed_up/fps
+    times = np.arange(t_start, t_end, time_per_frame)
+    
+    def update(time, ax, fig):
+        for ax_ in fig.get_axes():
+            ax_.clear()
+        ax = plot_function(ax=ax, t_start=t_start, t_end=time)
+        return ax
+    
+    ax = plot_function()
+    from matplotlib.animation import FuncAnimation
+    anim = FuncAnimation(plt.gcf(), update, frames=times, fargs=(ax, plt.gcf(),), interval=1000/fps)        
+    plt.close()
+    return anim
+
+def plot_trajectory(
+        trajectory,
+        time_stamps,
+        ax=None, 
+        t_start=None, 
+        t_end=None,
+        **plot_kwargs):
+    """Plots a trajectory over a given time range
+
+    Params
+    ------
+    ax: matplotlib axis, optional
+        Axis to plot on
+    t_start: float, optional
+        Start time
+    t_end: float, optional
+        End time
+    trajectory: np.ndarray, optional
+        Trajectory to plot, shape (T, 2)
+    time_stamps: np.ndarray, optional
+        Time stamps for the trajectory, shape (T,)
+    plot_kwargs: dict, optional
+        Additional plotting arguments like 'color', 'alpha', 'label', 'scatter_points', 'show_line', 'linewidth', 'title', 'xlabel', 'ylabel'
+
+    Returns
+    -------
+    ax: matplotlib axis
+    """
+    if t_start is None: t_start = time_stamps[0]
+    if t_end is None: t_end = time_stamps[-1]
+    if ax is None: 
+        fig, ax = plt.subplots(1, 1, figsize=(6, 6)) 
+    
+    id_start, id_end = np.argmin(np.abs(time_stamps - t_start)), np.argmin(np.abs(time_stamps-t_end))
+    trajectory_ = trajectory[id_start:id_end]
+
+    # Get plot kwargs
+    color = plot_kwargs.get('color', 'k')
+    scatter_points = plot_kwargs.get('scatter_points', True)
+    show_line = plot_kwargs.get('show_line', True)
+    linewidth = plot_kwargs.get('linewidth',1)
+    title = plot_kwargs.get('title',None)
+    xlabel = plot_kwargs.get('xlabel','x [m]')
+    ylabel = plot_kwargs.get('ylabel','y [m]')
+    alpha = plot_kwargs.get('alpha',1)
+    label = plot_kwargs.get('label',None)
+    min_x = plot_kwargs.get('min_x',trajectory[:,0].min().round(1))
+    max_x = plot_kwargs.get('max_x',trajectory[:,0].max().round(1))
+    min_y = plot_kwargs.get('min_y',trajectory[:,1].min().round(1))
+    max_y = plot_kwargs.get('max_y',trajectory[:,1].max().round(1))
+    
+    if show_line:
+        ax.plot(trajectory_[:,0],trajectory_[:,1],color=color, linewidth=linewidth, alpha=alpha, label=label)
+    if scatter_points:
+        ax.scatter(trajectory_[:,0],trajectory_[:,1],color=color, linewidth=0, s=6, alpha=alpha)
+    ax.set_xlim(min_x, max_x); ax.set_ylim(min_y, max_y); ax.set_aspect('equal', 'box')
+    ax.set_xticks([min_x, max_x]); ax.set_yticks([min_y, max_y])
+    ax.set_xlabel(xlabel)
+    ax.set_ylabel(ylabel)
+    if title is not None: ax.set_title(title)
+
+    return ax 
+
+def plot_ellipse(ax, mean, cov, color):
+    lambda_, v = np.linalg.eig(cov)
+    lambda_ = np.sqrt(lambda_) # convert from variance to standard deviation along the eigenvectors
+    ell = matplotlib.patches.Ellipse(xy=mean,
+                                     width=lambda_[0]*2, 
+                                     height=lambda_[1]*2,
+                                     angle=np.rad2deg(np.arctan(v[:, 0][1] / v[:, 0][0])),
+                                     lw=1, 
+                                     fill=True, 
+                                     edgecolor=color,
+                                     facecolor=color,
+                                     alpha=0.5,)
+    ax.add_artist(ell)
+    return ax