diff --git a/CITATION.cff b/CITATION.cff
index 7d05622..c693726 100644
--- a/CITATION.cff
+++ b/CITATION.cff
@@ -8,8 +8,8 @@ message: 'To cite package "animovement" in publications use:'
 type: software
 license: MIT
 title: 'animovement: An R toolbox for analysing animal movement across space and time.'
-version: 0.5.0
-date-released: 2024-12-04
+version: 0.5.1
+date-released: 2024-12-07
 abstract: An R toolbox for analysing animal movement across space and time.
 authors:
 - family-names: Roald-Arbøl
diff --git a/DESCRIPTION b/DESCRIPTION
index 616a110..78d26ff 100644
--- a/DESCRIPTION
+++ b/DESCRIPTION
@@ -1,7 +1,7 @@
 Package: animovement
 Type: Package
 Title: An R toolbox for analysing animal movement across space and time
-Version: 0.5.0
+Version: 0.5.1
 Authors@R: 
     person(
       "Mikkel", 
diff --git a/NAMESPACE b/NAMESPACE
index a60bbcf..1c8a610 100644
--- a/NAMESPACE
+++ b/NAMESPACE
@@ -24,7 +24,7 @@ export(read_treadmill)
 export(read_trex)
 export(set_start_datetime)
 export(set_uuid)
-export(smooth_track)
+export(smooth_movement)
 export(translate_coords)
 export(validate_animalta)
 export(validate_deeplabcut_csv)
diff --git a/NEWS.md b/NEWS.md
index 8d33bdb..294f9ba 100644
--- a/NEWS.md
+++ b/NEWS.md
@@ -1,3 +1,11 @@
+# animovement 0.5.1
+Previously, `smooth_track` was only built to smoothen trackball data. In this fix, we fixed it by doing the following:
+- Changed the name to `smooth_movement`
+- Added parameter `use_derivatives` (default: `FALSE`) which is how trackball data should be smoothed
+- Written the non-derivative code
+
+So now `smooth_movement` should be able to smooth all your movement data!
+
 # animovement 0.5.0
 A big update! There are three major updates:
 - We finally stabilised on a **data format** and **implemented it for all readers**! 
diff --git a/R/calculate_statistics.R b/R/calculate_statistics.R
index ca587dd..ee2787f 100644
--- a/R/calculate_statistics.R
+++ b/R/calculate_statistics.R
@@ -22,6 +22,10 @@ calculate_statistics <- function(
     straightness = c("A", "B", "C", "D")) {
   validate_statistics()
 
+  # Ungroup (makes summarise work better)
+  data <- data |>
+    dplyr::ungroup()
+
   # Calculate translational and rotational separately (maybe?) and gather at the end
   totals <- data |>
     dplyr::summarise(across(c("distance", "rotation"), ~ collapse::fsum(abs(.x)), .names = "total_{.col}"),
diff --git a/R/clean_movement.R b/R/clean_movement.R
index f5a35d8..416737a 100644
--- a/R/clean_movement.R
+++ b/R/clean_movement.R
@@ -1,4 +1,4 @@
-#' Smooth tracks
+#' Smooth movement
 #'
 #' @description
 #' `r lifecycle::badge('experimental')`
@@ -11,17 +11,60 @@
 #' @param data Data frame
 #' @param method Which smoothing method to use. options:  "rolling_median (default), "rolling_mean".
 #' @param window_width How many observations to use for rolling window filters (e.g. "rolling_mean" or "rolling_median").
+#' @param use_derivatives whether to use the derivatives (difference between frames) to perform the smoothing. Useful for trackball data.
 #'
 #' @return A movement data frame
 #' @export
 #' @import dplyr
 #' @importFrom roll roll_mean roll_median
 #'
-smooth_track <- function(
+smooth_movement <- function(
     data,
     method = c("rolling_median"),
-    window_width = 5) {
+    window_width = 5,
+    use_derivatives = FALSE) {
+
+  # Quick checks on the data
+  ensure_output_header_names(data)
+  ensure_output_header_class(data)
+
+  data <- data |>
+    dplyr::group_by(.data$individual, .data$keypoint)
+
   # Back-transform to dx and dy
+  if (use_derivatives == TRUE){
+    data <- smooth_derivatives(data, method, window_width)
+  }
+
+  # Rolling mean
+  else if (method == "rolling_mean"){
+    data <- data |>
+      dplyr::mutate(
+        x = roll::roll_mean(.data$x, width = window_width),
+        y = roll::roll_mean(.data$y, width = window_width)
+      )
+  }
+
+  # Rolling median
+  else if (method == "rolling_median") {
+    data <- data |>
+      dplyr::mutate(
+        x = roll::roll_median(.data$x, width = window_width),
+        y = roll::roll_median(.data$y, width = window_width)
+      )
+  }
+
+  return(data)
+}
+
+#' Smooth derivatives (when original data measured differences between movements)
+#'
+#' @inheritParams smooth_movement
+#'
+#' @keywords internal
+smooth_derivatives <- function(data, method, window_width){
+
+  # Back-transform to derivatives
   data <- data |>
     dplyr::mutate(
       dx = .data$x - lag(.data$x),
@@ -49,6 +92,7 @@ smooth_track <- function(
       x = cumsum(coalesce(.data$dx, 0)) + .data$dx * 0,
       y = cumsum(coalesce(.data$dy, 0)) + .data$dy * 0
     ) |>
-    select(-.data$dx, -.data$dy)
+    select(-"dx", -"dy")
+
   return(data)
 }
diff --git a/R/validator_output.R b/R/validator_output.R
index 931faa5..e731b4c 100644
--- a/R/validator_output.R
+++ b/R/validator_output.R
@@ -15,6 +15,7 @@ ensure_output_header_class <- function(
   data <- data |>
     dplyr::select(all_of(expected_headers))
   header_classes <- sapply(data, class)
+  header_classes <- sapply(header_classes, function(x) if_else(x == "integer", "numeric", x))
   if (!all(expected_header_class == header_classes)){
     cli::cli_abort("Expected output headers to be {expected_header_class}, got {header_classes}.")
   }
diff --git a/_pkgdown.yml b/_pkgdown.yml
index 2c224c7..39add60 100644
--- a/_pkgdown.yml
+++ b/_pkgdown.yml
@@ -29,7 +29,7 @@ reference:
 - title: "Clean"
   desc: These functions ensure that your data is ready for analysis.
   contents:
-  - smooth_track
+  - smooth_movement
   - clean_kinematics
 
 - title: "Transformations"
diff --git a/inst/CITATION b/inst/CITATION
index fe29ed9..3d788ef 100644
--- a/inst/CITATION
+++ b/inst/CITATION
@@ -9,4 +9,4 @@ bibentry(bibtype = "Misc",
          year = "2024",
          url = "http://www.roald-arboel.com/animovement/",
          abstract = "An R toolbox for analysing animal movement across space and time.",
-         version = "0.5.0")
+         version = "0.5.1")
diff --git a/man/smooth_derivatives.Rd b/man/smooth_derivatives.Rd
new file mode 100644
index 0000000..a475010
--- /dev/null
+++ b/man/smooth_derivatives.Rd
@@ -0,0 +1,19 @@
+% Generated by roxygen2: do not edit by hand
+% Please edit documentation in R/clean_movement.R
+\name{smooth_derivatives}
+\alias{smooth_derivatives}
+\title{Smooth derivatives (when original data measured differences between movements)}
+\usage{
+smooth_derivatives(data, method, window_width)
+}
+\arguments{
+\item{data}{Data frame}
+
+\item{method}{Which smoothing method to use. options:  "rolling_median (default), "rolling_mean".}
+
+\item{window_width}{How many observations to use for rolling window filters (e.g. "rolling_mean" or "rolling_median").}
+}
+\description{
+Smooth derivatives (when original data measured differences between movements)
+}
+\keyword{internal}
diff --git a/man/smooth_track.Rd b/man/smooth_movement.Rd
similarity index 72%
rename from man/smooth_track.Rd
rename to man/smooth_movement.Rd
index 9ee0bd3..b0d47f4 100644
--- a/man/smooth_track.Rd
+++ b/man/smooth_movement.Rd
@@ -1,10 +1,15 @@
 % Generated by roxygen2: do not edit by hand
 % Please edit documentation in R/clean_movement.R
-\name{smooth_track}
-\alias{smooth_track}
-\title{Smooth tracks}
+\name{smooth_movement}
+\alias{smooth_movement}
+\title{Smooth movement}
 \usage{
-smooth_track(data, method = c("rolling_median"), window_width = 5)
+smooth_movement(
+  data,
+  method = c("rolling_median"),
+  window_width = 5,
+  use_derivatives = FALSE
+)
 }
 \arguments{
 \item{data}{Data frame}
@@ -12,6 +17,8 @@ smooth_track(data, method = c("rolling_median"), window_width = 5)
 \item{method}{Which smoothing method to use. options:  "rolling_median (default), "rolling_mean".}
 
 \item{window_width}{How many observations to use for rolling window filters (e.g. "rolling_mean" or "rolling_median").}
+
+\item{use_derivatives}{whether to use the derivatives (difference between frames) to perform the smoothing. Useful for trackball data.}
 }
 \value{
 A movement data frame
diff --git a/tests/testthat/test-output.R b/tests/testthat/test-output.R
index 5cb9f0d..8d121c0 100644
--- a/tests/testthat/test-output.R
+++ b/tests/testthat/test-output.R
@@ -63,12 +63,12 @@ test_that("Test output header classes", {
   expect_no_error(ensure_output_header_class(df_bonsai))
   expect_no_error(ensure_output_header_class(df_dlc_single))
   expect_no_error(ensure_output_header_class(df_dlc_multi))
-  expect_no_error(ensure_output_header_class(df_idtracker_h5, expected_header_class = c("integer", "factor", "factor", "numeric", "numeric", "numeric")))
+  expect_no_error(ensure_output_header_class(df_idtracker_h5))
   expect_no_error(ensure_output_header_class(df_idtracker_csv))
   expect_no_error(ensure_output_header_class(df_lightningpose_single))
   expect_no_error(ensure_output_header_class(df_lightningpose_twoview))
-  expect_no_error(ensure_output_header_class(df_sleap_single, expected_header_class = c("integer", "factor", "factor", "numeric", "numeric", "numeric")))
-  expect_no_error(ensure_output_header_class(df_sleap_multi, expected_header_class = c("integer", "factor", "factor", "numeric", "numeric", "numeric")))
+  expect_no_error(ensure_output_header_class(df_sleap_single))
+  expect_no_error(ensure_output_header_class(df_sleap_multi))
   expect_no_error(ensure_output_header_class(df_trex))
   expect_no_error(ensure_output_header_class(df_trackball))
 })
diff --git a/tests/testthat/test-smooth.R b/tests/testthat/test-smooth.R
new file mode 100644
index 0000000..89267a2
--- /dev/null
+++ b/tests/testthat/test-smooth.R
@@ -0,0 +1,57 @@
+# Test arguments
+
+# Paths to test data
+path_animalta_raw <- here::here("tests", "data", "animalta", "single_individual_multi_arena.csv")
+path_animalta_detailed <- here::here("tests", "data", "animalta", "variable_individuals_single_arena.csv")
+path_bonsai <- here::here("tests", "data", "bonsai", "LI850.csv")
+path_dlc_single <- here::here("tests", "data", "deeplabcut", "mouse_single.csv")
+path_dlc_multi <- here::here("tests", "data", "deeplabcut", "mouse_multi.csv")
+path_idtracker_h5 <- here::here("tests", "data", "idtrackerai", "trajectories.h5")
+path_idtracker_csv <- here::here("tests", "data", "idtrackerai", "trajectories_csv", "trajectories.csv")
+path_idtracker_csv_probabilities <- here::here("tests", "data", "idtrackerai", "trajectories_csv", "id_probabilities.csv")
+path_lightningpose_single <- here::here("tests", "data", "lightningpose", "mouse_single.csv")
+path_lightningpose_twoview <- here::here("tests", "data", "lightningpose", "mouse_twoview.csv")
+path_sleap_single <- here::here("tests", "data", "sleap", "SLEAP_single-mouse_EPM.analysis.h5")
+path_sleap_multi <- here::here("tests", "data", "sleap", "SLEAP_three-mice_Aeon_mixed-labels.analysis.h5")
+path_trex <- here::here("tests", "data", "trex", "beetle.csv")
+path_trackball_1 <- here::here("tests", "data", "single", "opticalflow_sensor_1.csv")
+path_trackball_2 <- here::here("tests", "data", "single", "opticalflow_sensor_2.csv")
+
+# Read data snippets
+df_animalta_raw <- read_animalta(path_animalta_raw, detailed = FALSE)
+df_animalta_detailed <- read_animalta(path_animalta_detailed, detailed = TRUE)
+df_bonsai <- read_bonsai(path_bonsai)
+df_dlc_single <- read_deeplabcut(path_dlc_single)
+df_dlc_multi <- read_deeplabcut(path_dlc_multi)
+df_idtracker_h5 <- read_idtracker(path_idtracker_h5)
+df_idtracker_csv <- read_idtracker(path_idtracker_csv, path_idtracker_csv_probabilities)
+df_lightningpose_single <- read_lightningpose(path_lightningpose_single)
+df_lightningpose_twoview <- read_lightningpose(path_lightningpose_twoview)
+df_sleap_single <- read_sleap(path_sleap_single)
+df_sleap_multi <- read_sleap(path_sleap_multi)
+df_trex <- read_trex(path_trex)
+df_trackball <- read_trackball(
+  paths = c(path_trackball_1, path_trackball_2),
+  setup = "of_free",
+  sampling_rate = 60,
+  col_time = 4,
+  distance_scale = 394,
+  distance_unit = NULL
+)
+
+# Check that smoothing functions work
+test_that("Test output header names", {
+  expect_no_error(smooth_movement(df_animalta_raw))
+  expect_no_error(smooth_movement(df_animalta_detailed))
+  expect_no_error(smooth_movement(df_bonsai))
+  expect_no_error(smooth_movement(df_dlc_single))
+  expect_no_error(smooth_movement(df_dlc_multi))
+  expect_no_error(smooth_movement(df_idtracker_h5))
+  expect_no_error(smooth_movement(df_idtracker_csv))
+  expect_no_error(smooth_movement(df_lightningpose_single))
+  expect_no_error(smooth_movement(df_lightningpose_twoview))
+  expect_no_error(smooth_movement(df_sleap_single))
+  expect_no_error(smooth_movement(df_sleap_multi))
+  expect_no_error(smooth_movement(df_trex))
+  expect_no_error(smooth_movement(df_trackball, use_derivatives = TRUE))
+})
diff --git a/vignettes/articles/calculate-kinematics.Rmd b/vignettes/articles/calculate-kinematics.Rmd
index 4338d15..c09efbf 100644
--- a/vignettes/articles/calculate-kinematics.Rmd
+++ b/vignettes/articles/calculate-kinematics.Rmd
@@ -13,6 +13,7 @@ knitr::opts_chunk$set(
 library(animovement)
 library(tibble)
 library(ggplot2)
+library(tidyr)
 library(dplyr, warn.conflicts = FALSE)
 library(readxl)
 library(here)
@@ -34,7 +35,7 @@ df <- read_trackball(
 
 # Smooth tracks
 df_smooth <- df |>
-  smooth_track(method = "rolling_median", window_width = 5)
+  smooth_movement(method = "rolling_median", window_width = 5, use_derivatives = TRUE)
 ```
 
 ## Calculate kinematics
@@ -102,7 +103,7 @@ f <- df_kinematics_clean |>
 
 library(performance)
 df_kinematics_clean |>
-  na.omit() |>
+  tidyr::drop_na(rotation) |>
   select(rotation) |>
   performance::check_outliers()
 ```
diff --git a/vignettes/articles/calculate-summary-statistics.Rmd b/vignettes/articles/calculate-summary-statistics.Rmd
index 07fb532..5657332 100644
--- a/vignettes/articles/calculate-summary-statistics.Rmd
+++ b/vignettes/articles/calculate-summary-statistics.Rmd
@@ -34,7 +34,7 @@ df <- read_trackball(
 
 # Smooth tracks
 df_kinematics_clean <- df |>
-  smooth_track(method = "rolling_median", window_width = 5) |>
+  smooth_movement(method = "rolling_median", window_width = 5, use_derivatives = TRUE) |>
   calculate_kinematics() |>
   clean_kinematics()
 ```
diff --git a/vignettes/articles/clean-tracks.Rmd b/vignettes/articles/clean-tracks.Rmd
index 61e473d..fdaa26d 100644
--- a/vignettes/articles/clean-tracks.Rmd
+++ b/vignettes/articles/clean-tracks.Rmd
@@ -125,7 +125,7 @@ Let's try smoothing our data with a `rolling_mean` filter with 0.5 second (30 ob
 
 ```{r}
 df_smooth <- df |>
-  smooth_track(method = "rolling_median", window_width = 5)
+  smooth_movement(method = "rolling_median", window_width = 5, use_derivatives = TRUE)
 ```
 
 Let's visualise how they compare. Note that although the difference may seem negligible when plotting paths, they may become important when computing derivatives such as velocity and acceleration.