diff --git a/.github/workflows/full_tests.yml b/.github/workflows/full_tests.yml
new file mode 100644
index 0000000..4906926
--- /dev/null
+++ b/.github/workflows/full_tests.yml
@@ -0,0 +1,26 @@
+name: Test on Ubuntu
+
+on:
+  pull_request:
+      branches: [dev]
+      types: [synchronize, opened, reopened]
+
+
+jobs:
+  build-and-test:
+
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v3
+    - name: Set up Python
+      uses: actions/setup-python@v4
+      with:
+        python-version: "3.10"
+    - name: Install package
+      run: |
+        python -m pip install --upgrade pip
+        pip install .[test]
+    - name: Pytest
+      run: |
+        pytest -v
diff --git a/.github/workflows/publish-to-pypi.yml b/.github/workflows/publish-to-pypi.yml
new file mode 100644
index 0000000..7d5c266
--- /dev/null
+++ b/.github/workflows/publish-to-pypi.yml
@@ -0,0 +1,28 @@
+name: Release to PyPI
+
+on:
+  push:
+    tags:
+       - '*'
+jobs:
+  release:
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v2
+    - name: Set up Python 3.10
+      uses: actions/setup-python@v4
+      with:
+        python-version: "3.10"
+    - name: Install Tools
+      run: |
+        python -m pip install --upgrade pip
+        pip install setuptools wheel twine build
+    - name: Package and Upload
+      env:
+        STACKMANAGER_VERSION: ${{ github.event.release.tag_name }}
+        TWINE_USERNAME: __token__
+        TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }}
+      run: |
+        python -m build --sdist --wheel
+        twine upload dist/*
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
new file mode 100644
index 0000000..7be81c7
--- /dev/null
+++ b/.pre-commit-config.yaml
@@ -0,0 +1,36 @@
+repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.5.0
+    hooks:
+      - id: fix-encoding-pragma
+        exclude: tests/test_data
+      - id: trailing-whitespace
+        exclude: tests/test_data
+      - id: end-of-file-fixer
+        exclude: tests/test_data
+      - id: check-docstring-first
+      - id: debug-statements
+      - id: check-toml
+      - id: check-yaml
+        exclude: tests/test_data
+      - id: requirements-txt-fixer
+      - id: detect-private-key
+      - id: check-merge-conflict
+
+  - repo: https://github.com/psf/black
+    rev: 24.4.2
+    hooks:
+      - id: black
+        exclude: tests/test_data
+      - id: black-jupyter
+
+  - repo: https://github.com/pycqa/isort
+    rev: 5.13.2
+    hooks:
+      - id: isort
+        args: ["--profile", "black"]
+
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.4.4
+    hooks:
+      - id: ruff
diff --git a/examples/tracking_plot.py b/examples/tracking_plot.py
index 7bc421d..ec7649e 100644
--- a/examples/tracking_plot.py
+++ b/examples/tracking_plot.py
@@ -5,10 +5,21 @@
 from scipy.ndimage.measurements import center_of_mass
 
 
-def plot_path(x, y, t, box_size, spike_times=None,
-              color='grey', alpha=0.5, origin='upper',
-              spike_color='r', rate_markersize=False, markersize=10.,
-              animate=False, ax=None):
+def plot_path(
+    x,
+    y,
+    t,
+    box_size,
+    spike_times=None,
+    color="grey",
+    alpha=0.5,
+    origin="upper",
+    spike_color="r",
+    rate_markersize=False,
+    markersize=10.0,
+    animate=False,
+    ax=None,
+):
     """
     Plot path visited
 
@@ -39,8 +50,7 @@ def plot_path(x, y, t, box_size, spike_times=None,
     """
     if ax is None:
         fig = plt.figure()
-        ax = fig.add_subplot(
-            111, xlim=[0, box_size], ylim=[0, box_size], aspect=1)
+        ax = fig.add_subplot(111, xlim=[0, box_size], ylim=[0, box_size], aspect=1)
 
     ax.plot(x, y, c=color, alpha=alpha)
     if spike_times is not None:
@@ -49,20 +59,36 @@ def plot_path(x, y, t, box_size, spike_times=None,
 
         if rate_markersize:
             markersize = spikes_in_bin[is_spikes_in_bin] * markersize
-        ax.scatter(x[:-1][is_spikes_in_bin], y[:-1][is_spikes_in_bin],
-                   facecolor=spike_color, edgecolor=spike_color,
-                   s=markersize)
+        ax.scatter(
+            x[:-1][is_spikes_in_bin],
+            y[:-1][is_spikes_in_bin],
+            facecolor=spike_color,
+            edgecolor=spike_color,
+            s=markersize,
+        )
 
     ax.grid(False)
-    if origin == 'upper':
+    if origin == "upper":
         ax.invert_yaxis()
     return ax
 
 
-def animate_path(x, y, t, box_size, spike_times=None,
-              color='grey', alpha=0.5, origin='upper',
-              spike_color='r', rate_markersize=False, markersize=10.,
-              animate=False, ax=None, title=''):
+def animate_path(
+    x,
+    y,
+    t,
+    box_size,
+    spike_times=None,
+    color="grey",
+    alpha=0.5,
+    origin="upper",
+    spike_color="r",
+    rate_markersize=False,
+    markersize=10.0,
+    animate=False,
+    ax=None,
+    title="",
+):
     """
     Plot path visited
 
@@ -93,35 +119,35 @@ def animate_path(x, y, t, box_size, spike_times=None,
     """
     if ax is None:
         fig = plt.figure()
-        ax = fig.add_subplot(
-            111, xlim=[0, box_size], ylim=[0, box_size], aspect=1)
+        ax = fig.add_subplot(111, xlim=[0, box_size], ylim=[0, box_size], aspect=1)
     if spike_times is not None:
         spikes_in_bin, _ = np.histogram(spike_times, t)
         is_spikes_in_bin = np.array(spikes_in_bin, dtype=bool)
 
         if rate_markersize:
-            markersizes = spikes_in_bin[is_spikes_in_bin]*markersize
+            markersizes = spikes_in_bin[is_spikes_in_bin] * markersize
         else:
-            markersizes = markersize*np.ones(is_spikes_in_bin.size)
+            markersizes = markersize * np.ones(is_spikes_in_bin.size)
     ax.set_title(title)
     ax.grid(False)
-    if origin == 'upper':
+    if origin == "upper":
         ax.invert_yaxis()
     import time
+
     plt.show()
     for idx, x, y, active, msize in zip(range(len(x)), x, y):
         ax.plot(x, y, c=color, alpha=alpha)
         if spike_times is not None:
             if is_spikes_in_bin[idx]:
-                ax.scatter(x, y, facecolor=spike_color, edgecolor=spike_color,
-                           s=markersizes[idx])
+                ax.scatter(x, y, facecolor=spike_color, edgecolor=spike_color, s=markersizes[idx])
         time.sleep(0.1)  # plt.pause(0.0001)
         plt.draw()
     return ax
 
 
-def plot_head_direction_rate(spike_times, ang_bins, rate_in_ang, projection='polar',
-                             normalization=False, ax=None, color='k'):
+def plot_head_direction_rate(
+    spike_times, ang_bins, rate_in_ang, projection="polar", normalization=False, ax=None, color="k"
+):
     """
 
 
@@ -142,8 +168,9 @@ def plot_head_direction_rate(spike_times, ang_bins, rate_in_ang, projection='pol
     out : ax
     """
     import math
+
     if normalization:
-        rate_in_ang = normalize(rate_in_ang, mode='minmax')
+        rate_in_ang = normalize(rate_in_ang, mode="minmax")
     if ax is None:
         fig = plt.figure()
         ax = fig.add_subplot(111, projection=projection)
@@ -151,7 +178,7 @@ def plot_head_direction_rate(spike_times, ang_bins, rate_in_ang, projection='pol
     if projection is None:
         ax.set_xticks(range(0, 360 + 60, 60))
         ax.set_xlim(0, 360)
-    elif projection == 'polar':
+    elif projection == "polar":
         ang_bins = [math.radians(deg) for deg in ang_bins]
         bin_size = math.radians(bin_size)
         ax.set_xticks([0, np.pi])
@@ -159,9 +186,20 @@ def plot_head_direction_rate(spike_times, ang_bins, rate_in_ang, projection='pol
     return ax
 
 
-def plot_ratemap(x, y, t, spike_times, bin_size=0.05, box_size=1,
-                 box_size=1, vmin=0, ax=None, smoothing=.05,
-                 origin='upper', cmap='viridis'):
+def plot_ratemap(
+    x,
+    y,
+    t,
+    spike_times,
+    bin_size=0.05,
+    box_size=1,
+    box_size=1,
+    vmin=0,
+    ax=None,
+    smoothing=0.05,
+    origin="upper",
+    cmap="viridis",
+):
     """
 
 
@@ -184,19 +222,17 @@ def plot_ratemap(x, y, t, spike_times, bin_size=0.05, box_size=1,
         fig = plt.figure()
         ax = fig.add_subplot(111, xlim=[0, 1], ylim=[0, 1], aspect=1)
 
-    map = SpatialMap(
-        x, y, t, spike_times, bin_size=bin_size, box_size=box_size)
+    map = SpatialMap(x, y, t, spike_times, bin_size=bin_size, box_size=box_size)
     rate_map = map.rate_map(smoothing)
-    ax.imshow(rate_map, interpolation='none', origin=origin,
-              extent=(0, 1, 0, 1), vmin=vmin, cmap=cmap)
-    ax.set_title('%.2f Hz' % np.nanmax(rate_map))
+    ax.imshow(rate_map, interpolation="none", origin=origin, extent=(0, 1, 0, 1), vmin=vmin, cmap=cmap)
+    ax.set_title("%.2f Hz" % np.nanmax(rate_map))
     ax.grid(False)
     return ax
 
 
-def plot_occupancy(x, y, t, bin_size=0.05, box_size=1, box_size=1,
-                  vmin=0, ax=None, convolve=True,
-                  origin='upper', cmap='jet'):
+def plot_occupancy(
+    x, y, t, bin_size=0.05, box_size=1, box_size=1, vmin=0, ax=None, convolve=True, origin="upper", cmap="jet"
+):
     """
 
 
@@ -219,10 +255,10 @@ def plot_occupancy(x, y, t, bin_size=0.05, box_size=1, box_size=1,
         fig = plt.figure()
         ax = fig.add_subplot(111, xlim=[0, 1], ylim=[0, 1], aspect=1)
 
-    occ_map = occupancy_map(x, y, t, bin_size=bin_size, box_size=box_size,
-                             box_size=box_size, convolve=convolve)
-    cax = ax.imshow(occ_map, interpolation='none', origin=origin,
-                   extent=(0, 1, 0, 1), vmin=vmin, cmap=cmap, aspect='auto')
+    occ_map = occupancy_map(x, y, t, bin_size=bin_size, box_size=box_size, box_size=box_size, convolve=convolve)
+    cax = ax.imshow(
+        occ_map, interpolation="none", origin=origin, extent=(0, 1, 0, 1), vmin=vmin, cmap=cmap, aspect="auto"
+    )
     # ax.set_title('%.2f s' % np.nanmax(occ_map))
     ax.grid(False)
     return cax, np.nanmax(occ_map)
diff --git a/pyproject.toml b/pyproject.toml
new file mode 100644
index 0000000..64376bd
--- /dev/null
+++ b/pyproject.toml
@@ -0,0 +1,58 @@
+[project]
+name = "spatial_maps"
+version = "0.2.0"
+authors = [
+    { name = "Mikkel Lepperod", email = "mikkel@simula.no" },
+    { name = "Alessio Buccino", email = "alessiop.buccino@gmail.com" },
+]
+
+description = "Compute spatial maps for neural data."
+readme = "README.md"
+requires-python = ">=3.10"
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+]
+
+dependencies = [
+    "numpy<2",
+    "scipy",
+    "scikit-image",
+    "astropy",
+    "pandas",
+    "elephant",
+    "matplotlib"
+]
+
+[project.urls]
+homepage = "https://github.com/CINPLA/spatial-maps"
+repository = "https://github.com/CINPLA/spatial-maps"
+
+[build-system]
+requires = ["setuptools>=62.0"]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools]
+include-package-data = true
+
+[tool.setuptools.packages.find]
+where = ["src"]
+include = ["spatial_maps*"]
+namespaces = false
+
+[project.optional-dependencies]
+dev = ["pre-commit", "black[jupyter]", "isort", "ruff"]
+test = ["pytest", "pytest-cov", "pytest-dependency", "mountainsort5"]
+docs = ["sphinx-gallery", "sphinx_rtd_theme"]
+full = [
+    "spatial_maps[dev]",
+    "spatial_maps[test]",
+    "spatial_maps[docs]",
+]
+
+[tool.coverage.run]
+omit = ["tests/*"]
+
+[tool.black]
+line-length = 120
diff --git a/setup.py b/setup.py
index be6606d..7a58127 100644
--- a/setup.py
+++ b/setup.py
@@ -1,31 +1,6 @@
 # -*- coding: utf-8 -*-
-from setuptools import setup
-import os
 
-from setuptools import setup, find_packages
+import setuptools
 
-long_description = open("README.md").read()
-
-install_requires = [
-    'numpy>=1.9',
-    'scipy',
-    'astropy',
-    'pandas>=0.14.1',
-    'elephant',
-    'matplotlib']
-extras_require = {
-    'testing': ['pytest'],
-    'docs': ['numpydoc>=0.5',
-             'sphinx>=1.2.2',
-             'sphinx_rtd_theme']
-}
-
-setup(
-    name="spatial_maps",
-    install_requires=install_requires,
-    tests_require=install_requires,
-    extras_require=extras_require,
-    packages=find_packages(),
-    include_package_data=True,
-    version='0.1',
-)
+if __name__ == "__main__":
+    setuptools.setup()
diff --git a/spatial_maps/__init__.py b/spatial_maps/__init__.py
deleted file mode 100644
index c26fa44..0000000
--- a/spatial_maps/__init__.py
+++ /dev/null
@@ -1,11 +0,0 @@
-from .maps import SpatialMap
-from .gridcells import (
-    gridness, spacing_and_orientation, separate_fields_by_distance)
-from .fields import (
-    calculate_field_centers, separate_fields_by_laplace,
-    find_peaks, which_field, compute_crossings)
-from .bordercells import border_score
-from .stats import (
-    sparsity, selectivity, information_rate, information_specificity,
-    prob_dist)
-from .tools import autocorrelation, fftcorrelate2d, nancorrelate2d
diff --git a/spatial_maps/tests/test_maps.py b/spatial_maps/tests/test_maps.py
deleted file mode 100644
index bfa474f..0000000
--- a/spatial_maps/tests/test_maps.py
+++ /dev/null
@@ -1,70 +0,0 @@
-import numpy as np
-import pytest
-from spatial_maps.maps import SpatialMap
-import quantities as pq
-from spatial_maps.tools import make_test_grid_rate_map, make_test_spike_map
-
-
-def test_rate_map():
-    box_size = [1., 1.]
-    rate = 5.
-    bin_size = [.01, .01]
-    n_step=10**4
-    step_size=.1
-    sigma=0.1
-    spacing=0.3
-    smoothing = .03
-
-    rate_map_true, pos_fields, xbins, ybins = make_test_grid_rate_map(
-        sigma=sigma, spacing=spacing, amplitude=rate, box_size=box_size,
-        bin_size=bin_size)
-
-    x, y, t, spikes = make_test_spike_map(
-        pos_fields=pos_fields, box_size=box_size, rate=rate,
-        n_step=n_step, step_size=step_size, sigma=sigma)
-    smap = SpatialMap(
-        x, y, t, spikes, box_size=box_size, bin_size=bin_size)
-    rate_map = smap.rate_map(smoothing)
-
-    diff = rate_map_true - rate_map
-    X, Y = np.meshgrid(xbins, ybins)
-
-    samples = []
-    for p in pos_fields:
-        mask = np.sqrt((X - p[0])**2 + (Y - p[1])**2) < .1
-        samples.append(diff[mask])
-    peak_diff = np.abs(np.mean([s.min() for s in samples if s.size > 0]))
-    assert peak_diff < 0.5
-
-
-def test_spatial_rate_map_diag():
-    N = 10
-    bin_size = 1
-    box_size = 1.0
-    x = np.linspace(0., box_size, N)
-    y = np.linspace(0., box_size, N)
-    t = np.linspace(0.1, 10.1, N)
-    spike_times = np.arange(0.1, 10.1, .5)
-    map = SpatialMap(
-        x, y, t, spike_times, box_size=box_size, bin_size=bin_size)
-    ratemap = map.rate_map(0)
-    print(ratemap)
-    assert all(np.diff(np.diag(ratemap)) < 1e-10)
-    assert ratemap.shape == (int(box_size / bin_size), int(box_size / bin_size))
-
-
-def test_occupancy_map_diag():
-    N = 3
-    bin_size = .5
-    box_size = 1.5
-    x = np.linspace(0., box_size, N)
-    y = np.linspace(0., box_size, N)
-    t = np.linspace(0, 10., N)
-
-    map = SpatialMap(
-        x, y, t, [], box_size=box_size, bin_size=bin_size)
-    occmap_expected = np.array([[5, 0, 0],
-                                [0, 5, 0],
-                                [0, 0, 5]])
-    occmap = map.occupancy_map(0)
-    assert np.array_equal(occmap, occmap_expected)
diff --git a/spatial_maps/tests/test_stats.py b/spatial_maps/tests/test_stats.py
deleted file mode 100644
index 2ea3e49..0000000
--- a/spatial_maps/tests/test_stats.py
+++ /dev/null
@@ -1,36 +0,0 @@
-import numpy as np
-import pytest
-
-
-def test_calc_population_vector_correlation():
-    from spatial_maps.stats import population_vector_correlation as pvcorr
-    rmaps1 = np.array([
-        [
-            [1, 0.1],
-            [0.1, 4]
-        ],
-        [
-            [6, 0.1],
-            [0.1, 2]
-        ],
-        [
-            [2, 0.1],
-            [0.1, 3]
-        ]])
-    rmaps2 = np.array([
-        [
-            [2, 0.2],
-            [0.2, 8]
-        ],
-        [
-            [12, 0.2],
-            [0.2, 4]
-        ],
-        [
-            [4, 0.2],
-            [0.2, 6]
-        ]])
-    rmaps2 += 10e-5
-    pv = pvcorr(rmaps1, rmaps2)
-    err = pv-1
-    assert err < 10e-5
diff --git a/src/spatial_maps/__init__.py b/src/spatial_maps/__init__.py
new file mode 100644
index 0000000..a45b503
--- /dev/null
+++ b/src/spatial_maps/__init__.py
@@ -0,0 +1,6 @@
+from .maps import SpatialMap
+from .gridcells import gridness, spacing_and_orientation, separate_fields_by_distance
+from .fields import calculate_field_centers, separate_fields_by_laplace, find_peaks, which_field, compute_crossings
+from .bordercells import border_score
+from .stats import sparsity, selectivity, information_rate, information_specificity, prob_dist
+from .tools import autocorrelation, fftcorrelate2d, nancorrelate2d
diff --git a/spatial_maps/bordercells.py b/src/spatial_maps/bordercells.py
similarity index 85%
rename from spatial_maps/bordercells.py
rename to src/spatial_maps/bordercells.py
index c4144cf..bfb823c 100644
--- a/spatial_maps/bordercells.py
+++ b/src/spatial_maps/bordercells.py
@@ -26,9 +26,9 @@ def border_score(rate_map, fields):
     if np.all(fields == 0):
         raise ValueError("Must have at least one field")
 
-    inner = np.zeros(np.array(rate_map.shape)-(4,4),dtype=bool)
-    wall = np.pad(inner, 1, 'constant', constant_values=[[0,0],[1,0]])
-    wall = np.pad(wall, 1, 'constant', constant_values=0)
+    inner = np.zeros(np.array(rate_map.shape) - (4, 4), dtype=bool)
+    wall = np.pad(inner, 1, "constant", constant_values=[[0, 0], [1, 0]])
+    wall = np.pad(wall, 1, "constant", constant_values=0)
     max_extent = 0
 
     ones = np.ones_like(rate_map)
@@ -36,8 +36,8 @@ def border_score(rate_map, fields):
     for i in range(4):
         borders = np.logical_and(fields > 0, wall)
         extents = labeled_comprehension(
-            input=borders, labels=fields, index=None, func=np.sum,
-            out_dtype=np.int64, default=0)
+            input=borders, labels=fields, index=None, func=np.sum, out_dtype=np.int64, default=0
+        )
         max_extent = np.max([max_extent, np.max(extents)])
 
         # dont rotate the fourth time
@@ -47,7 +47,7 @@ def border_score(rate_map, fields):
 
     x = np.linspace(-0.5, 0.5, rate_map.shape[1])
     y = np.linspace(-0.5, 0.5, rate_map.shape[0])
-    X,Y = np.meshgrid(x,y)
+    X, Y = np.meshgrid(x, y)
 
     # create linear increasing value towards middle
     dist_to_nearest_wall = 1 - (np.abs(X + Y) + np.abs(X - Y))
diff --git a/spatial_maps/fields.py b/src/spatial_maps/fields.py
similarity index 86%
rename from spatial_maps/fields.py
rename to src/spatial_maps/fields.py
index ac0cd1f..e4d4bd2 100644
--- a/spatial_maps/fields.py
+++ b/src/spatial_maps/fields.py
@@ -4,6 +4,7 @@
 from scipy.interpolate import interp2d, interp1d
 from .tools import fftcorrelate2d, autocorrelation
 
+
 def border_score(rate_map, fields):
     raise DeprecationWarning('This function is moved to "spatial_maps.bordercells"')
     return spatial_maps.stats(rate_map, fields)
@@ -17,12 +18,13 @@ def find_peaks(image):
     peaks : array
         coordinates for peaks in image as [row, column]
     """
+    from scipy.ndimage import maximum_filter
     image = image.copy()
     image[~np.isfinite(image)] = 0
-    image_max = filters.maximum_filter(image, 3)
-    is_maxima = (image == image_max)
+    image_max = maximum_filter(image, 3)
+    is_maxima = image == image_max
     labels, num_objects = ndimage.label(is_maxima)
-    indices = np.arange(1, num_objects+1)
+    indices = np.arange(1, num_objects + 1)
     peaks = ndimage.maximum_position(image, labels=labels, index=indices)
     peaks = np.array(peaks)
     center = (np.array(image.shape) - 1) / 2
@@ -32,7 +34,7 @@ def find_peaks(image):
 
 
 def sort_fields_by_rate(rate_map, fields, func=None):
-    '''Sort fields by the rate value of each field
+    """Sort fields by the rate value of each field
     Parameters
     ----------
     rate_map : array
@@ -44,12 +46,11 @@ def sort_fields_by_rate(rate_map, fields, func=None):
     -------
     sorted_fields : array
         Sorted fields
-    '''
+    """
     indx = np.sort(np.unique(fields.ravel()))
     func = func or np.max
     # Sort by largest peak
-    rate_means = ndimage.labeled_comprehension(
-        rate_map, fields, indx, func, np.float64, 0)
+    rate_means = ndimage.labeled_comprehension(rate_map, fields, indx, func, np.float64, 0)
     sort = np.argsort(rate_means)[::-1]
 
     # new rate map with fields > min_size, sorted
@@ -61,7 +62,7 @@ def sort_fields_by_rate(rate_map, fields, func=None):
 
 
 def remove_fields_by_area(fields, minimum_field_area):
-    '''Sets fields below minimum area to zero, measured as the number of bins in a field.
+    """Sets fields below minimum area to zero, measured as the number of bins in a field.
     Parameters
     ----------
     fields : array
@@ -72,7 +73,7 @@ def remove_fields_by_area(fields, minimum_field_area):
     -------
     fields
         Fields with number of bins below minimum_field_area are set to zero
-    '''
+    """
     if not isinstance(minimum_field_area, (int, np.integer)):
         raise ValueError("'minimum_field_area' should be int")
 
@@ -83,7 +84,7 @@ def remove_fields_by_area(fields, minimum_field_area):
     # fields[fields_area < minimum_field_area] = 0
 
     labels, counts = np.unique(fields, return_counts=True)
-    for (lab, count) in zip(labels, counts):
+    for lab, count in zip(labels, counts):
         if lab != 0:
             if count < minimum_field_area:
                 fields[fields == lab] = 0
@@ -150,8 +151,9 @@ def separate_fields_by_dilation(rate_map, seed=2.5, sigma=2.5, minimum_field_are
     see https://scikit-image.org/docs/stable/auto_examples/color_exposure/plot_regional_maxima.html
     """
     from skimage.morphology import reconstruction
+
     rate_map_norm = (rate_map - rate_map.mean()) / rate_map.std()
-    dilated = reconstruction(rate_map_norm - seed, rate_map_norm, method='dilation')
+    dilated = reconstruction(rate_map_norm - seed, rate_map_norm, method="dilation")
     rate_map_reconstructed = rate_map_norm - dilated
 
     l = ndimage.gaussian_laplace(rate_map_reconstructed, sigma)
@@ -194,28 +196,26 @@ def separate_fields_by_laplace_of_gaussian(rate_map, sigma=2, minimum_field_area
     return fields
 
 
-def calculate_field_centers(rate_map, labels, center_method='maxima'):
+def calculate_field_centers(rate_map, labels, center_method="maxima"):
     """Finds center of fields at labels.
     :Authors:
         Halvard Sutterud <halvard.sutterud@gmail.com>
     """
 
     from scipy import ndimage
+
     indices = np.arange(1, np.max(labels) + 1)
-    if center_method == 'maxima':
-        bc = ndimage.maximum_position(
-            rate_map, labels=labels, index=indices)
-    elif center_method == 'center_of_mass':
-        bc = ndimage.center_of_mass(
-            rate_map, labels=labels, index=indices)
+    if center_method == "maxima":
+        bc = ndimage.maximum_position(rate_map, labels=labels, index=indices)
+    elif center_method == "center_of_mass":
+        bc = ndimage.center_of_mass(rate_map, labels=labels, index=indices)
     else:
-        raise ValueError(
-            "invalid center_method flag '{}'".format(center_method))
+        raise ValueError("invalid center_method flag '{}'".format(center_method))
     if not bc:
         # empty list
         return bc
     bc = np.array(bc)
-    bc[:,[0, 1]] = bc[:,[1, 0]] # y, x -> x, y
+    bc[:, [0, 1]] = bc[:, [1, 0]]  # y, x -> x, y
     return bc
 
 
@@ -238,13 +238,13 @@ def which_field(x, y, fields, box_size):
         arraylike x and y with fields-labeled indices
     """
 
-    if len(x)!= len(y):
-        raise ValueError('x and y must have same length')
+    if len(x) != len(y):
+        raise ValueError("x and y must have same length")
 
     sx, sy = fields.shape
     # bin sizes
-    dx = box_size[0]/sx
-    dy = box_size[1]/sy
+    dx = box_size[0] / sx
+    dy = box_size[1] / sy
     x_bins = dx + np.arange(0, box_size[0] + dx, dx)
     y_bins = dy + np.arange(0, box_size[1] + dx, dy)
     # x_bins = np.arange(0, box_size[0] + dx, dx)
@@ -253,9 +253,9 @@ def which_field(x, y, fields, box_size):
     iy = np.digitize(y, y_bins)
 
     # fix for boundaries:
-    ix[ix==sx] = sx-1
-    iy[iy==sy] = sy-1
-    return np.array(fields[ix,iy])
+    ix[ix == sx] = sx - 1
+    iy[iy == sy] = sy - 1
+    return np.array(fields[ix, iy])
 
 
 def compute_crossings(field_indices):
@@ -270,13 +270,13 @@ def compute_crossings(field_indices):
     """
     # make sure to start and end outside fields
     field_indices = np.concatenate(([0], field_indices.astype(bool).astype(int), [0]))
-    enter, = np.where(np.diff(field_indices) == 1)
-    exit, = np.where(np.diff(field_indices) == -1)
+    (enter,) = np.where(np.diff(field_indices) == 1)
+    (exit,) = np.where(np.diff(field_indices) == -1)
     assert len(enter) == len(exit), (len(enter), len(exit))
     return enter, exit
 
 
-def distance_to_edge_function(x_c, y_c, field, box_size, interpolation='linear'):
+def distance_to_edge_function(x_c, y_c, field, box_size, interpolation="linear"):
     """Returns a function which for a given angle returns the distance to
     the edge of the field from the center.
     Parameters:
@@ -287,6 +287,7 @@ def distance_to_edge_function(x_c, y_c, field, box_size, interpolation='linear')
     """
 
     from skimage import measure
+
     contours = measure.find_contours(field, 0.8)
 
     box_dim = np.array(box_size)
@@ -300,13 +301,13 @@ def distance_to_edge_function(x_c, y_c, field, box_size, interpolation='linear')
     edge_y = edge_y[a_sort]
 
     # # Fill in edge values for the interpolation
-    pad_a = np.pad(angles, 2, mode='linear_ramp', end_values=(0, 2 * np.pi))
+    pad_a = np.pad(angles, 2, mode="linear_ramp", end_values=(0, 2 * np.pi))
     ev_x = (edge_x[0] + edge_x[-1]) / 2
-    pad_x = np.pad(edge_x, 2, mode='linear_ramp', end_values=ev_x)
+    pad_x = np.pad(edge_x, 2, mode="linear_ramp", end_values=ev_x)
     ev_y = (edge_y[0] + edge_y[-1]) / 2
-    pad_y = np.pad(edge_y, 2, mode='linear_ramp', end_values=ev_y)
+    pad_y = np.pad(edge_y, 2, mode="linear_ramp", end_values=ev_y)
 
-    if interpolation=='cubic':
+    if interpolation == "cubic":
         mask = np.where(np.diff(pad_a) == 0)
         pad_a = np.delete(pad_a, mask)
         pad_x = np.delete(pad_x, mask)
@@ -318,7 +319,7 @@ def distance_to_edge_function(x_c, y_c, field, box_size, interpolation='linear')
     def dist_func(angle):
         x = x_func(angle)
         y = y_func(angle)
-        dist = np.sqrt((x - x_c)**2 + (y - y_c)**2)
+        dist = np.sqrt((x - x_c) ** 2 + (y - y_c) ** 2)
         return dist
 
     return dist_func
@@ -360,10 +361,8 @@ def map_pass_to_unit_circle(x, y, t, x_c, y_c, field=None, box_size=None, dist_f
         placecell firing in open environment
     """
     if dist_func is None:
-        assert field is not None and box_size is not None, (
-            'either provide "dist_func" or "field" and "box_size"')
-        dist_func= distance_to_edge_function(
-            x_c, y_c, field, box_size, interpolation='linear')
+        assert field is not None and box_size is not None, 'either provide "dist_func" or "field" and "box_size"'
+        dist_func = distance_to_edge_function(x_c, y_c, field, box_size, interpolation="linear")
     pos = np.array((x, y))
 
     # vector from pos to center p
@@ -387,9 +386,9 @@ def map_pass_to_unit_circle(x, y, t, x_c, y_c, field=None, box_size=None, dist_f
     # is toward positive x
     theta = (angle - np.arctan2(mean_velocity[1], mean_velocity[0])) % (2 * np.pi)
 
-    w_pdcd = (angle - np.arctan2(velocity[1], velocity[0]))
+    w_pdcd = angle - np.arctan2(velocity[1], velocity[0])
     pdcd = r * np.cos(w_pdcd)
 
-    w_pdmd = (angle - np.arctan2(mean_velocity[1], mean_velocity[0]))
+    w_pdmd = angle - np.arctan2(mean_velocity[1], mean_velocity[0])
     pdmd = r * np.cos(w_pdmd)
     return r, theta, pdcd, pdmd
diff --git a/spatial_maps/gridcells.py b/src/spatial_maps/gridcells.py
similarity index 90%
rename from spatial_maps/gridcells.py
rename to src/spatial_maps/gridcells.py
index cce1408..15cec22 100644
--- a/spatial_maps/gridcells.py
+++ b/src/spatial_maps/gridcells.py
@@ -34,7 +34,7 @@ def separate_fields_by_distance(rate_map, factor=0.7):
     """
     import scipy.spatial as spatial
 
-    acorr = autocorrelation(rate_map, mode='full', normalize=True)
+    acorr = autocorrelation(rate_map, mode="full", normalize=True)
     acorr_maxima = find_peaks(acorr)
 
     def place_field_radius(auto_correlation, maxima):
@@ -42,7 +42,7 @@ def place_field_radius(auto_correlation, maxima):
         center = map_size / 2
         distances = np.linalg.norm(maxima - center, axis=1)
         distances_sorted = sorted(distances)
-        min_distance = distances_sorted[1] # the first one is basically the center
+        min_distance = distances_sorted[1]  # the first one is basically the center
         return factor * min_distance / 2
 
     # TODO verify this for an example where there are fields too close
@@ -50,7 +50,7 @@ def too_close_removed(rate_map, rate_map_maxima, place_field_radius):
         result = []
         rate_map_maxima_value = rate_map[tuple(rate_map_maxima.T)]
         distances = spatial.distance.cdist(rate_map_maxima, rate_map_maxima)
-        too_close_pairs = np.where(distances < place_field_radius*2)
+        too_close_pairs = np.where(distances < place_field_radius * 2)
         not_accepted = []
 
         for i, j in zip(*too_close_pairs):
@@ -99,9 +99,10 @@ def peak_to_peak_distance(sorted_peaks, index_a, index_b):
 
 def rotate_corr(acorr, mask):
     import numpy.ma as ma
-    from scipy.ndimage.interpolation import rotate
+    from scipy.ndimage import rotate
+
     m_acorr = ma.masked_array(acorr, mask=mask)
-    angles = range(30, 180+30, 30)
+    angles = range(30, 180 + 30, 30)
     corr = []
     # Rotate and compute correlation coefficient
     for angle in angles:
@@ -114,7 +115,7 @@ def rotate_corr(acorr, mask):
 
 
 def gridness(rate_map, return_mask=False):
-    '''
+    """
     Calculates gridness based on the autocorrelation of a rate map.
     The Pearson's product-moment correlation coefficients are calculated between A and A_r,
     where A_r is the rotated version of A at 30, 60, 90, 120, and 150 degrees.
@@ -133,11 +134,12 @@ def gridness(rate_map, return_mask=False):
     Returns
     -------
     out : gridness
-    '''
+    """
     import numpy.ma as ma
+
     rate_map = rate_map.copy()
     rate_map[~np.isfinite(rate_map)] = 0
-    acorr = autocorrelation(rate_map, mode='full', normalize=True)
+    acorr = autocorrelation(rate_map, mode="full", normalize=True)
 
     acorr_maxima = find_peaks(acorr)
     inner_radius = 0.5 * peak_to_peak_distance(acorr_maxima, 0, 1)
@@ -151,12 +153,12 @@ def gridness(rate_map, return_mask=False):
     center = np.array(acorr.shape) / 2
     lower = (center - outer_radius).astype(int)
     upper = (center + outer_radius).astype(int)
-    acorr = acorr[lower[0]:upper[0], lower[1]:upper[1]]
+    acorr = acorr[lower[0] : upper[0], lower[1] : upper[1]]
 
     # create a mask
     ylen, xlen = acorr.shape  # ylen, xlen is the correct order for meshgrid
-    x = np.linspace(- xlen / 2., xlen / 2., xlen)
-    y = np.linspace(- ylen / 2., ylen / 2., ylen)
+    x = np.linspace(-xlen / 2.0, xlen / 2.0, xlen)
+    y = np.linspace(-ylen / 2.0, ylen / 2.0, ylen)
     X, Y = np.meshgrid(x, y)
     distance_map = np.sqrt(X**2 + Y**2)
     mask = (distance_map < inner_radius) | (distance_map > outer_radius)
@@ -217,7 +219,7 @@ def spacing_and_orientation(peaks, box_size):
     spacing = np.mean(closest_distances)
 
     # sort by angle
-    a = np.arctan2(closest_relpos[:,0], closest_relpos[:,1]) % (2 * np.pi)
+    a = np.arctan2(closest_relpos[:, 0], closest_relpos[:, 1]) % (2 * np.pi)
     a_sort = np.argsort(a)
 
     # extract lowest angle in radians
@@ -226,12 +228,10 @@ def spacing_and_orientation(peaks, box_size):
     return spacing, orientation
 
 
-def autocorrelation_centers(rate_map, threshold=0, center_method='maxima'):
+def autocorrelation_centers(rate_map, threshold=0, center_method="maxima"):
     # autocorrelate. Returns array (2x - 1) the size of rate_map
-    acorr = fftcorrelate2d(
-        rate_map, rate_map, mode='full', normalize=True)
+    acorr = fftcorrelate2d(rate_map, rate_map, mode="full", normalize=True)
     fields = separate_fields_by_laplace(rate_map, threshold=threshold)
 
-    field_centers = calculate_field_centers(
-        rate_map, fields, center_method=center_method)
+    field_centers = calculate_field_centers(rate_map, fields, center_method=center_method)
     return field_centers
diff --git a/spatial_maps/maps.py b/src/spatial_maps/maps.py
similarity index 83%
rename from spatial_maps/maps.py
rename to src/spatial_maps/maps.py
index 8dd7a1c..5e36a5c 100644
--- a/spatial_maps/maps.py
+++ b/src/spatial_maps/maps.py
@@ -69,15 +69,11 @@ def interpolate_nan_2D(array, method="nearest"):
     y1 = yy[~array.mask]
     newarr = array[~array.mask]
 
-    return interpolate.griddata(
-        (x1, y1), newarr.ravel(), (xx, yy), method=method, fill_value=0
-    )
+    return interpolate.griddata((x1, y1), newarr.ravel(), (xx, yy), method=method, fill_value=0)
 
 
 class SpatialMap:
-    def __init__(
-        self, smoothing=0.05, box_size=[1.0, 1.0], bin_size=0.02, bin_count=None
-    ):
+    def __init__(self, smoothing=0.05, box_size=[1.0, 1.0], bin_size=0.02, bin_count=None):
         """
         Parameters
         ----------
@@ -99,11 +95,7 @@ def __init__(
 
     def spike_map(self, x, y, t, spike_times, mask_zero_occupancy=True, **kwargs):
         spmap = _spike_map(x, y, t, spike_times, self.xbins, self.ybins)
-        spmap = (
-            smooth_map(spmap, self.bin_size, self.smoothing, **kwargs)
-            if self.smoothing
-            else spmap
-        )
+        spmap = smooth_map(spmap, self.bin_size, self.smoothing, **kwargs) if self.smoothing else spmap
         if mask_zero_occupancy:
             spmap[_occupancy_map(x, y, t, self.xbins, self.ybins) == 0] = np.nan
         return spmap
@@ -111,25 +103,12 @@ def spike_map(self, x, y, t, spike_times, mask_zero_occupancy=True, **kwargs):
     def occupancy_map(self, x, y, t, mask_zero_occupancy=True, **kwargs):
         ocmap = _occupancy_map(x, y, t, self.xbins, self.ybins)
         ocmap_copy = copy(ocmap)  # to mask zero occupancy after smoothing
-        ocmap = (
-            smooth_map(ocmap, self.bin_size, self.smoothing, **kwargs)
-            if self.smoothing
-            else ocmap
-        )
+        ocmap = smooth_map(ocmap, self.bin_size, self.smoothing, **kwargs) if self.smoothing else ocmap
         if mask_zero_occupancy:
             ocmap[ocmap_copy == 0] = np.nan
         return ocmap
 
-    def rate_map(
-        self,
-        x,
-        y,
-        t,
-        spike_times,
-        mask_zero_occupancy=True,
-        interpolate_invalid=False,
-        **kwargs
-    ):
+    def rate_map(self, x, y, t, spike_times, mask_zero_occupancy=True, interpolate_invalid=False, **kwargs):
         """Calculate rate map as spike_map / occupancy_map
         Parameters
         ----------
@@ -144,9 +123,7 @@ def rate_map(
         -------
         rate_map : array
         """
-        spike_map = self.spike_map(
-            x, y, t, spike_times, mask_zero_occupancy=mask_zero_occupancy, **kwargs
-        )
+        spike_map = self.spike_map(x, y, t, spike_times, mask_zero_occupancy=mask_zero_occupancy, **kwargs)
         # to avoid infinity (x/0) we set zero occupancy to nan
         occupancy_map = self.occupancy_map(x, y, t, mask_zero_occupancy=True, **kwargs)
         rate_map = spike_map / occupancy_map
diff --git a/spatial_maps/stats.py b/src/spatial_maps/stats.py
similarity index 92%
rename from spatial_maps/stats.py
rename to src/spatial_maps/stats.py
index 25eddb7..9612594 100644
--- a/spatial_maps/stats.py
+++ b/src/spatial_maps/stats.py
@@ -2,20 +2,19 @@
 
 
 def _inf_rate(rate_map, px):
-    '''
+    """
     A helper function for information rate.
 
     Originally from https://github.com/MattNolanLab/gridcells
-    '''
+    """
     tmp_rate_map = rate_map.copy()
     tmp_rate_map[np.isnan(tmp_rate_map)] = 0
     avg_rate = np.sum(np.ravel(tmp_rate_map * px))
-    return (np.nansum(np.ravel(tmp_rate_map * np.log2(tmp_rate_map/avg_rate) *
-            px)), avg_rate)
+    return (np.nansum(np.ravel(tmp_rate_map * np.log2(tmp_rate_map / avg_rate) * px)), avg_rate)
 
 
 def sparsity(rate_map, px):
-    '''
+    """
     Compute sparsity of a rate map, The sparsity  measure is an adaptation
     to space. The adaptation measures the fraction of the environment  in which
     a cell is  active. A sparsity of, 0.1 means that the place field of the
@@ -36,7 +35,7 @@ def sparsity(rate_map, px):
     .. [2] Skaggs, W. E., McNaughton, B. L., Wilson, M., & Barnes, C. (1996).
        Theta phase precession in hippocampal neuronal populations and the
        compression of temporal sequences. Hippocampus, 6, 149-172.
-    '''
+    """
     tmp_rate_map = rate_map.copy()
     tmp_rate_map[np.isnan(tmp_rate_map)] = 0
     avg_rate = np.sum(np.ravel(tmp_rate_map * px))
@@ -45,7 +44,7 @@ def sparsity(rate_map, px):
 
 
 def selectivity(rate_map, px):
-    '''
+    """
     "The selectivity measure max(rate)/mean(rate)  of the cell. The more
     tightly concentrated  the cell's activity, the higher the selectivity.
     A cell with no spatial tuning at all will  have a  selectivity of 1" [2]_.
@@ -59,7 +58,7 @@ def selectivity(rate_map, px):
     -------
     out : float
         selectivity
-    '''
+    """
     tmp_rate_map = rate_map.copy()
     tmp_rate_map[np.isnan(tmp_rate_map)] = 0
     avg_rate = np.sum(np.ravel(tmp_rate_map * px))
@@ -68,7 +67,7 @@ def selectivity(rate_map, px):
 
 
 def information_rate(rate_map, px):
-    '''
+    """
     Compute information rate of a cell given variable x.
     A simple algorithm devised by [1]_. This computes the spatial information
     rate of cell spikes given variable x (e.g. position, head direction) in
@@ -111,12 +110,12 @@ def information_rate(rate_map, px):
     .. [1] Skaggs, W.E. et al., 1993. An Information-Theoretic Approach to
        Deciphering the Hippocampal Code. In Advances in Neural Information
        Processing Systems 5. pp. 1030-1037.
-    '''
+    """
     return _inf_rate(rate_map, px)[0]
 
 
 def information_specificity(rate_map, px):
-    '''
+    """
     Compute the 'specificity' of the cell firing rate to a variable X.
     Compute :func:`information_rate` for this cell and divide by the average
     firing rate of the cell. See [1]_ for more information.
@@ -135,13 +134,13 @@ def information_specificity(rate_map, px):
     -------
     I : float
         Information in bits/spike.
-    '''
+    """
     I, avg_rate = _inf_rate(rate_map, px)
     return I / avg_rate
 
 
 def prob_dist(x, y, bins):
-    '''
+    """
     Calculate a probability distribution for animal positions in an arena.
 
     Parameters
@@ -155,14 +154,14 @@ def prob_dist(x, y, bins):
     dist : numpy.ndarray
         Probability distribution for the positional data. The first dimension
         is the y axis, the second dimension is the x axis.
-    '''
+    """
 
     H, _, _ = np.histogram2d(x, y, bins=bins, density=False)
-    return (H / len(x))
+    return H / len(x)
 
 
 def prob_dist_1d(x, bins):
-    '''
+    """
     Calculate a probability distribution for animal positions in an arena.
 
     Parameters
@@ -175,15 +174,13 @@ def prob_dist_1d(x, bins):
     dist : numpy.ndarray
         Probability distribution for the positional data. The first dimension
         is the y axis, the second dimension is the x axis.
-    '''
+    """
 
     H, _ = np.histogram(x, bins=bins, density=False)
     return (H / len(x)).T
 
 
-def population_vector_correlation(rmaps1, rmaps2,
-                                  mask_nans=False,
-                                  return_corr_coeff_map=False):
+def population_vector_correlation(rmaps1, rmaps2, mask_nans=False, return_corr_coeff_map=False):
     """
     Calcualte population vector correlation between two
     stacks of rate maps.
@@ -224,16 +221,12 @@ def population_vector_correlation(rmaps1, rmaps2,
                 bool_nan_xy1 = np.isnan(xy1)
                 bool_nan_xy2 = np.isnan(xy2)
 
-                mask_invalid = np.logical_or(
-                    bool_nan_xy1,
-                    bool_nan_xy2)
+                mask_invalid = np.logical_or(bool_nan_xy1, bool_nan_xy2)
                 mask_valid = ~mask_invalid
                 xy1 = xy1[mask_valid]
                 xy2 = xy2[mask_valid]
 
-            corr_coeff_map[i, j] = np.corrcoef(
-                xy1,
-                xy2)[0, 1]
+            corr_coeff_map[i, j] = np.corrcoef(xy1, xy2)[0, 1]
 
     pop_vec_corr = np.nanmean(corr_coeff_map)
 
diff --git a/spatial_maps/tools.py b/src/spatial_maps/tools.py
similarity index 89%
rename from spatial_maps/tools.py
rename to src/spatial_maps/tools.py
index 389c967..db977a3 100644
--- a/spatial_maps/tools.py
+++ b/src/spatial_maps/tools.py
@@ -82,21 +82,15 @@ def nancorrelate2d(X, Y, mode="frobenius") -> np.ndarray:
     result = np.zeros(X.shape)
     for i in range(X.shape[0]):
         for j in range(X.shape[1]):
-            scope_i = slice(
-                max(0, i - X.shape[0] // 2), min(i + X.shape[0] // 2, X.shape[0])
-            )
-            scope_j = slice(
-                max(0, j - X.shape[1] // 2), min(j + X.shape[1] // 2, X.shape[1])
-            )
+            scope_i = slice(max(0, i - X.shape[0] // 2), min(i + X.shape[0] // 2, X.shape[0]))
+            scope_j = slice(max(0, j - X.shape[1] // 2), min(j + X.shape[1] // 2, X.shape[1]))
             if mode == "pearson":
                 result[i, j] = ma.corrcoef(
                     X[scope_i, scope_j].flatten(),
                     Y[scope_i, scope_j][::-1][:, ::-1].flatten(),
                 )[0, 1]
             elif mode == "frobenius":  # scaled (average) frobenius inner product
-                result[i, j] = (
-                    X[scope_i, scope_j] * Y[scope_i, scope_j][::-1][:, ::-1]
-                ).mean()
+                result[i, j] = (X[scope_i, scope_j] * Y[scope_i, scope_j][::-1][:, ::-1]).mean()
             else:
                 raise NotImplementedError("Method does not have mode={}".format(mode))
 
@@ -154,18 +148,10 @@ def gaussian2D(amp, x, y, xc, yc, s):
 def gaussian2D_asym(pos, amplitude, xc, yc, sigma_x, sigma_y, theta):
     x, y = pos
 
-    a = (np.cos(theta) ** 2) / (2 * sigma_x ** 2) + (np.sin(theta) ** 2) / (
-        2 * sigma_y ** 2
-    )
-    b = -(np.sin(2 * theta)) / (4 * sigma_x ** 2) + (np.sin(2 * theta)) / (
-        4 * sigma_y ** 2
-    )
-    c = (np.sin(theta) ** 2) / (2 * sigma_x ** 2) + (np.cos(theta) ** 2) / (
-        2 * sigma_y ** 2
-    )
-    g = amplitude * np.exp(
-        -(a * ((x - xc) ** 2) + 2 * b * (x - xc) * (y - yc) + c * ((y - yc) ** 2))
-    )
+    a = (np.cos(theta) ** 2) / (2 * sigma_x**2) + (np.sin(theta) ** 2) / (2 * sigma_y**2)
+    b = -(np.sin(2 * theta)) / (4 * sigma_x**2) + (np.sin(2 * theta)) / (4 * sigma_y**2)
+    c = (np.sin(theta) ** 2) / (2 * sigma_x**2) + (np.cos(theta) ** 2) / (2 * sigma_y**2)
+    g = amplitude * np.exp(-(a * ((x - xc) ** 2) + 2 * b * (x - xc) * (y - yc) + c * ((y - yc) ** 2)))
     return g.ravel()
 
 
@@ -305,13 +291,8 @@ def random_walk(box_size, step_size, n_step, sampling_rate, low_pass=5):
     boundaries = np.array([(0, box_size[0]), (0, box_size[1])])
     size = np.diff(boundaries, axis=1).ravel()
     # "simulation"
-    trajectory = np.cumsum(
-        directions[np.random.randint(0, 9, (n_step,))] * step_size, axis=0
-    )
-    x, y = (
-        np.abs((trajectory + start - boundaries[:, 0] + size) % (2 * size) - size)
-        + boundaries[:, 0]
-    ).T
+    trajectory = np.cumsum(directions[np.random.randint(0, 9, (n_step,))] * step_size, axis=0)
+    x, y = (np.abs((trajectory + start - boundaries[:, 0] + size) % (2 * size) - size) + boundaries[:, 0]).T
 
     b, a = ss.butter(N=1, Wn=low_pass * 2 / sampling_rate)
     # zero phase shift filter
@@ -324,9 +305,7 @@ def random_walk(box_size, step_size, n_step, sampling_rate, low_pass=5):
     return x, y
 
 
-def make_test_spike_map(
-    rate, sigma, pos_fields, box_size, n_step=10 ** 4, step_size=0.05
-):
+def make_test_spike_map(rate, sigma, pos_fields, box_size, n_step=10**4, step_size=0.05):
     from scipy.interpolate import interp1d
 
     def infield(pos, pos_fields):
diff --git a/spatial_maps/tests/test_bordercells.py b/tests/test_bordercells.py
similarity index 78%
rename from spatial_maps/tests/test_bordercells.py
rename to tests/test_bordercells.py
index 532f9e9..db82504 100644
--- a/spatial_maps/tests/test_bordercells.py
+++ b/tests/test_bordercells.py
@@ -6,14 +6,15 @@ def test_border_score():
     from spatial_maps.bordercells import border_score
     from spatial_maps.tools import make_test_border_map
     from spatial_maps.fields import separate_fields_by_laplace
-    box_size = [1., 1.]
-    rate = 1.
-    bin_size = [.01, .01]
+
+    box_size = [1.0, 1.0]
+    rate = 1.0
+    bin_size = [0.01, 0.01]
 
     rate_map, pos_true, xbins, ybins = make_test_border_map(
-        sigma=0.05, amplitude=rate, offset=0, box_size=box_size,
-        bin_size=bin_size)
+        sigma=0.05, amplitude=rate, offset=0, box_size=box_size, bin_size=bin_size
+    )
 
     labels = separate_fields_by_laplace(rate_map, threshold=0)
     bs = border_score(rate_map, labels)
-    assert round(bs, 2) == .32
+    assert round(bs, 2) == 0.32
diff --git a/spatial_maps/tests/test_fields.py b/tests/test_fields.py
similarity index 53%
rename from spatial_maps/tests/test_fields.py
rename to tests/test_fields.py
index f2292fa..92755aa 100644
--- a/spatial_maps/tests/test_fields.py
+++ b/tests/test_fields.py
@@ -3,79 +3,57 @@
 import quantities as pq
 from spatial_maps.tools import make_test_grid_rate_map, make_test_border_map
 from spatial_maps.fields import (
-    separate_fields_by_laplace, find_peaks, calculate_field_centers,
-    in_field, distance_to_edge_function, map_pass_to_unit_circle)
+    separate_fields_by_laplace,
+    find_peaks,
+    calculate_field_centers,
+    distance_to_edge_function,
+    map_pass_to_unit_circle,
+)
 
 
 def test_find_peaks():
-    box_size = np.array([1., 1.])
-    rate = 5.
-    bin_size = [.01, .01]
-    sigma=0.05
-    spacing=0.3
+    box_size = np.array([1.0, 1.0])
+    rate = 5.0
+    bin_size = [0.01, 0.01]
+    sigma = 0.05
+    spacing = 0.3
 
     rate_map, pos_fields, xbins, ybins = make_test_grid_rate_map(
-        sigma=sigma, spacing=spacing, amplitude=rate, offset=0, box_size=box_size,
-        bin_size=bin_size, repeat=0)
+        sigma=sigma, spacing=spacing, amplitude=rate, offset=0, box_size=box_size, bin_size=bin_size, repeat=0
+    )
     peaks = find_peaks(rate_map)
-    pos_peaks = np.array([xbins[peaks[:,1]], ybins[peaks[:,0]]]).T
+    pos_peaks = np.array([xbins[peaks[:, 1]], ybins[peaks[:, 0]]]).T
     print(pos_peaks)
-    assert all(
-        [np.isclose(p, pos_peaks, rtol=1e-3).prod(axis=1).any()
-         for p in pos_fields])
+    assert all([np.isclose(p, pos_peaks, rtol=1e-3).prod(axis=1).any() for p in pos_fields])
 
 
 def test_separate_fields_by_laplace():
-    box_size = [1., 1.]
-    rate = 1.
-    bin_size = [.01, .01]
-    sigma=0.05
-    spacing=0.3
+    box_size = [1.0, 1.0]
+    rate = 1.0
+    bin_size = [0.01, 0.01]
+    sigma = 0.05
+    spacing = 0.3
 
     rate_map, pos_true, xbins, ybins = make_test_grid_rate_map(
-        sigma=sigma, spacing=spacing, amplitude=rate, offset=0.1, box_size=box_size,
-        bin_size=bin_size, orientation=0.1)
+        sigma=sigma, spacing=spacing, amplitude=rate, offset=0.1, box_size=box_size, bin_size=bin_size, orientation=0.1
+    )
 
     labels = separate_fields_by_laplace(rate_map, threshold=0)
     peaks = calculate_field_centers(rate_map, labels)
-    bump_centers = np.array([xbins[peaks[:,0]], ybins[peaks[:,1]]])
+    bump_centers = np.array([xbins[peaks[:, 0]], ybins[peaks[:, 1]]])
     # The position of a 2D bin is defined to be its center
     for p in pos_true:
         assert np.isclose(p, pos_true).prod(axis=1).any()
 
 
-def test_in_field():
-    n_bins   = 10
-    box_size = [1, 1]
-    bin_size = box_size[0] / n_bins
-
-    fields     = np.zeros((n_bins, n_bins))
-    fields[:5] = 1
-    fields[7]  = 2
-
-    # pick out center of bins
-    x = np.arange(bin_size/2, box_size[0], bin_size)
-    y = box_size[0] / 2 * np.ones_like(x)
-    true_value = [1, 1, 1, 1, 1, 0, 0, 2, 0, 0]
-    assert np.all(in_field(x, y, fields, box_size) == true_value)
-    # test edges
-    x = np.array([0, 1])
-    y = np.array([0.5, 0.5])
-    fields[:] = 0
-    fields[0] = 1
-    fields[-1] = 2
-    assert np.all(in_field(x, y, fields,box_size) == [1, 2])
-
-
 def test_distance_to_edge_function():
-    n_bins   = 10
+    n_bins = 10
     box_size = [1, 1]
     bin_size = box_size[0] / n_bins
 
-    field    = np.zeros((n_bins, n_bins))
+    field = np.zeros((n_bins, n_bins))
     field[2:8, 2:8] = 1
-    d = distance_to_edge_function(
-        0.5, 0.5, field, box_size, interpolation='linear')
+    d = distance_to_edge_function(0.5, 0.5, field, box_size, interpolation="linear")
 
     # assert edges 3/10 of the box size from center
     for a in [i * np.pi / 2 for i in range(4)]:
@@ -92,13 +70,14 @@ def test_distance_to_edge_function():
 
     # Greens theorem
     area = 0.5 * np.sum(x * dy - y * dx)
-    exact_area = np.sum(field) / np.size(field) * box_size[0]**2
+    exact_area = np.sum(field) / np.size(field) * box_size[0] ** 2
 
     assert np.abs(area - exact_area) / exact_area < 0.05
 
+
 def test_map_pass_to_unit_circle():
 
-    dist_func = lambda theta : 1
+    dist_func = lambda theta: 1
     x_c, y_c = (0.5, 0.5)
     theta = np.linspace(np.pi, 2 * np.pi, 100)
     t = theta
diff --git a/spatial_maps/tests/test_gridcells.py b/tests/test_gridcells.py
similarity index 67%
rename from spatial_maps/tests/test_gridcells.py
rename to tests/test_gridcells.py
index 798e06d..4430db9 100644
--- a/spatial_maps/tests/test_gridcells.py
+++ b/tests/test_gridcells.py
@@ -2,36 +2,34 @@
 import pytest
 from spatial_maps import SpatialMap
 import quantities as pq
-from spatial_maps.tools import (
-    make_test_grid_rate_map, make_test_spike_map, autocorrelation)
+from spatial_maps.tools import make_test_grid_rate_map, make_test_spike_map, autocorrelation
 from spatial_maps.fields import find_peaks
-from spatial_maps.gridcells import (
-    gridness, spacing_and_orientation, separate_fields_by_distance)
+from spatial_maps.gridcells import gridness, spacing_and_orientation, separate_fields_by_distance
 
 
 def test_gridness():
-    box_size = np.array([1., 1.])
-    rate = 5.
-    bin_size = [.01, .01]
+    box_size = np.array([1.0, 1.0])
+    rate = 5.0
+    bin_size = [0.01, 0.01]
     spacing_true = 0.3
 
     rate_map, pos_fields, xbins, ybins = make_test_grid_rate_map(
-        sigma=0.05, spacing=spacing_true, amplitude=rate, offset=0, box_size=box_size,
-        bin_size=bin_size)
+        sigma=0.05, spacing=spacing_true, amplitude=rate, offset=0, box_size=box_size, bin_size=bin_size
+    )
 
     g = gridness(rate_map)
     assert round(g, 1) == 1.3
 
 
 def test_spacing_and_orientation_from_true_peaks():
-    box_size = np.array([1., 1.])
-    rate = 5.
-    bin_size = [.01, .01]
+    box_size = np.array([1.0, 1.0])
+    rate = 5.0
+    bin_size = [0.01, 0.01]
     spacing_true = 0.3
 
     rate_map, pos_fields, xbins, ybins = make_test_grid_rate_map(
-        sigma=0.05, spacing=spacing_true, amplitude=rate, offset=0, box_size=box_size,
-        bin_size=bin_size)
+        sigma=0.05, spacing=spacing_true, amplitude=rate, offset=0, box_size=box_size, bin_size=bin_size
+    )
 
     spacing, orientation = spacing_and_orientation(pos_fields, box_size)
     assert spacing == spacing_true
@@ -39,15 +37,21 @@ def test_spacing_and_orientation_from_true_peaks():
 
 
 def test_spacing_and_orientation_from_autocorr():
-    box_size = np.array([1., 1.])
-    rate = 5.
-    bin_size = [.01, .01]
+    box_size = np.array([1.0, 1.0])
+    rate = 5.0
+    bin_size = [0.01, 0.01]
     spacing_true = 0.3
-    orientation_true = .3
+    orientation_true = 0.3
 
     rate_map, pos_fields, xbins, ybins = make_test_grid_rate_map(
-        sigma=0.05, spacing=spacing_true, amplitude=rate, offset=0, box_size=box_size,
-        bin_size=bin_size, orientation=orientation_true)
+        sigma=0.05,
+        spacing=spacing_true,
+        amplitude=rate,
+        offset=0,
+        box_size=box_size,
+        bin_size=bin_size,
+        orientation=orientation_true,
+    )
     autocorrelogram = autocorrelation(rate_map)
     peaks = find_peaks(autocorrelogram)
     real_peaks = peaks * bin_size
@@ -58,16 +62,16 @@ def test_spacing_and_orientation_from_autocorr():
 
 
 def test_separate_fields_by_distance():
-    box_size = [1., 1.]
-    rate = 1.
-    bin_size = [.01, .01]
+    box_size = [1.0, 1.0]
+    rate = 1.0
+    bin_size = [0.01, 0.01]
 
     rate_map, pos_true, xbins, ybins = make_test_grid_rate_map(
-        sigma=0.05, spacing=0.3, amplitude=rate, offset=0, box_size=box_size,
-        bin_size=bin_size)
+        sigma=0.05, spacing=0.3, amplitude=rate, offset=0, box_size=box_size, bin_size=bin_size
+    )
 
     peaks, radius = separate_fields_by_distance(rate_map)
-    bump_centers = np.array([xbins[peaks[:,0]], ybins[peaks[:,1]]])
+    bump_centers = np.array([xbins[peaks[:, 0]], ybins[peaks[:, 1]]])
     # The position of a 2D bin is defined to be its center
     for p in pos_true:
         assert np.isclose(p, pos_true).prod(axis=1).any()
@@ -83,18 +87,15 @@ def test_separate_fields_by_distance_2():
     for field in fields:
         dY = Y - field[0]
         dX = X - field[1]
-        rate_map += np.exp(-1/2*(dY**2 + dX**2)/10)  # Gaussian-ish
+        rate_map += np.exp(-1 / 2 * (dY**2 + dX**2) / 10)  # Gaussian-ish
 
     # should be removed by the algorithm because they are lower and close to existing fields
-    noise_fields = [
-        [60, 52],
-        [45, 35]
-    ]
+    noise_fields = [[60, 52], [45, 35]]
 
     for field in noise_fields:
         dY = Y - field[0]
         dX = X - field[1]
-        rate_map += 0.5 * np.exp(-1/2*(dY**2 + dX**2)/10)  # Gaussian-ish
+        rate_map += 0.5 * np.exp(-1 / 2 * (dY**2 + dX**2) / 10)  # Gaussian-ish
 
     found_fields, radius = separate_fields_by_distance(rate_map)
 
diff --git a/tests/test_maps.py b/tests/test_maps.py
new file mode 100644
index 0000000..ea809e0
--- /dev/null
+++ b/tests/test_maps.py
@@ -0,0 +1,67 @@
+import numpy as np
+import pytest
+from spatial_maps.maps import SpatialMap
+import quantities as pq
+from spatial_maps.tools import make_test_grid_rate_map, make_test_spike_map
+
+
+def test_rate_map():
+    box_size = [1.0, 1.0]
+    rate = 5.0
+    bin_size = [0.01, 0.01]
+    n_step = 10**4
+    step_size = 0.1
+    sigma = 0.1
+    spacing = 0.3
+    smoothing = 0.03
+
+    rate_map_true, pos_fields, xbins, ybins = make_test_grid_rate_map(
+        sigma=sigma, spacing=spacing, amplitude=rate, box_size=box_size, bin_size=bin_size
+    )
+
+    x, y, t, spikes = make_test_spike_map(
+        pos_fields=pos_fields, box_size=box_size, rate=rate, n_step=n_step, step_size=step_size, sigma=sigma
+    )
+    smap = SpatialMap(smoothing=smoothing, box_size=box_size, bin_size=bin_size)
+    rate_map = smap.rate_map(x, y, t, spikes)
+
+    rate_map[np.isnan(rate_map)] = 0
+    diff = rate_map_true - rate_map
+    X, Y = np.meshgrid(xbins, ybins)
+
+    samples = []
+    for p in pos_fields:
+        mask = np.sqrt((X - p[0]) ** 2 + (Y - p[1]) ** 2) < 0.1
+        samples.append(diff[mask])
+    peak_diff = np.abs(np.mean([s.min() for s in samples if s.size > 0]))
+    assert peak_diff < 0.5
+
+
+def test_spatial_rate_map_diag():
+    N = 10
+    bin_size = 1
+    box_size = 1.0
+    x = np.linspace(0.0, box_size, N)
+    y = np.linspace(0.0, box_size, N)
+    t = np.linspace(0.1, 10.1, N)
+    spike_times = np.arange(0.1, 10.1, 0.5)
+    sm = SpatialMap(box_size=box_size, bin_size=bin_size)
+    ratemap = sm.rate_map(x, y, t, spike_times)
+    print(ratemap)
+    assert all(np.diff(np.diag(ratemap)) < 1e-10)
+    assert ratemap.shape == (int(box_size / bin_size), int(box_size / bin_size))
+
+
+def test_occupancy_map_diag():
+    N = 3
+    bin_size = 0.5
+    box_size = 1.5
+    x = np.linspace(0.0, box_size, N)
+    y = np.linspace(0.0, box_size, N)
+    t = np.linspace(0, 10.0, N)
+
+    sm = SpatialMap(box_size=box_size, bin_size=bin_size)
+    occmap_expected = np.array([[5, 0, 0], [0, 5, 0], [0, 0, 5]])
+    occmap = sm.occupancy_map(x, y, t)
+    occmap[np.isnan(occmap)] = 0
+    assert np.array_equal(occmap, occmap_expected)
diff --git a/tests/test_stats.py b/tests/test_stats.py
new file mode 100644
index 0000000..a64be72
--- /dev/null
+++ b/tests/test_stats.py
@@ -0,0 +1,13 @@
+import numpy as np
+import pytest
+
+
+def test_calc_population_vector_correlation():
+    from spatial_maps.stats import population_vector_correlation as pvcorr
+
+    rmaps1 = np.array([[[1, 0.1], [0.1, 4]], [[6, 0.1], [0.1, 2]], [[2, 0.1], [0.1, 3]]])
+    rmaps2 = np.array([[[2, 0.2], [0.2, 8]], [[12, 0.2], [0.2, 4]], [[4, 0.2], [0.2, 6]]])
+    rmaps2 += 10e-5
+    pv = pvcorr(rmaps1, rmaps2)
+    err = pv - 1
+    assert err < 10e-5