diff --git a/src/reboost/hpge/psd.py b/src/reboost/hpge/psd.py
index 1254fc5..2b8f8cc 100644
--- a/src/reboost/hpge/psd.py
+++ b/src/reboost/hpge/psd.py
@@ -8,7 +8,7 @@
 import numpy as np
 import pint
 import pyg4ometry
-from lgdo import Array, VectorOfVectors
+from lgdo import Array
 from numpy.typing import ArrayLike, NDArray
 
 from .. import units
@@ -18,7 +18,7 @@
 log = logging.getLogger(__name__)
 
 
-def r90(edep: ak.Array, xloc: ak.Array, yloc: ak.Array, zloc: ak.Array) -> Array:
+def r90(edep: ak.Array, xloc: ak.Array, yloc: ak.Array, zloc: ak.Array) -> ak.Array:
     """R90 HPGe pulse shape heuristic.
 
     Parameters
@@ -31,18 +31,20 @@ def r90(edep: ak.Array, xloc: ak.Array, yloc: ak.Array, zloc: ak.Array) -> Array
         array of y coordinate position.
     zloc
         array of z coordinate position.
+
+    Returns
+    -------
+    calculated r90 for each hit
     """
+    pos = [units.units_conv_ak(pos, "mm") for pos in [xloc, yloc, zloc]]
+
     tot_energy = ak.sum(edep, axis=-1, keepdims=True)
 
     def eweight_mean(field, energy):
         return ak.sum(energy * field, axis=-1, keepdims=True) / tot_energy
 
     # Compute distance of each edep to the weighted mean
-    dist = np.sqrt(
-        (xloc - eweight_mean(edep, xloc)) ** 2
-        + (yloc - eweight_mean(edep, yloc)) ** 2
-        + (zloc - eweight_mean(edep, zloc)) ** 2
-    )
+    dist = np.sqrt(sum((p - eweight_mean(edep, p)) ** 2 for p in pos))
 
     # Sort distances and corresponding edep within each event
     sorted_indices = ak.argsort(dist, axis=-1)
@@ -76,16 +78,16 @@ def _ak_cumsum(layout, **_kwargs):
     r90_indices = ak.argmax(cumsum_edep_corrected >= threshold, axis=-1, keepdims=True)
     r90 = sorted_dist[r90_indices]
 
-    return Array(ak.flatten(r90).to_numpy())
+    return ak.Array(ak.flatten(r90), attrs={"units": "mm"})
 
 
 def drift_time(
-    xloc: ArrayLike,
-    yloc: ArrayLike,
-    zloc: ArrayLike,
+    xloc: ak.Array,
+    yloc: ak.Array,
+    zloc: ak.Array,
     dt_map: HPGeScalarRZField,
     coord_offset: pint.Quantity | pyg4ometry.gdml.Position = (0, 0, 0) * u.m,
-) -> VectorOfVectors:
+) -> ak.Array:
     """Calculates drift times for each step (cluster) in an HPGe detector.
 
     Parameters
@@ -106,7 +108,7 @@ def drift_time(
     # sanitize coord_offset
     coord_offset = units.pg4_to_pint(coord_offset)
 
-    # unit handling (for matching with drift time map units)
+    # unit handling (.r_units) for data in (xloc, yloc)]
     xu, yu = [units.units_convfact(data, dt_map.r_units) for data in (xloc, yloc)]
     zu = units.units_convfact(zloc, dt_map.z_units)
 
@@ -122,7 +124,6 @@ def _ak_dt_map(layouts, **_kwargs):
 
         return None
 
-    # transform coordinates
     xloc = xu * xloc - coord_offset[0].to(dt_map.r_units).m
     yloc = yu * yloc - coord_offset[1].to(dt_map.r_units).m
     zloc = zu * zloc - coord_offset[2].to(dt_map.z_units).m
@@ -133,16 +134,13 @@ def _ak_dt_map(layouts, **_kwargs):
         np.sqrt(xloc**2 + yloc**2),
         zloc,
     )
-    return VectorOfVectors(
-        dt_values,
-        attrs={"units": units.unit_to_lh5_attr(dt_map.φ_units)},
-    )
+    return units.attach_units(ak.Array(dt_values), units.unit_to_lh5_attr(dt_map.φ_units))
 
 
 def drift_time_heuristic(
-    drift_time: ArrayLike,
-    edep: ArrayLike,
-) -> Array:
+    drift_time: ak.Array,
+    edep: ak.Array,
+) -> ak.Array:
     """HPGe drift-time-based pulse-shape heuristic.
 
     See :func:`_drift_time_heuristic_impl` for a description of the algorithm.
@@ -160,11 +158,12 @@ def drift_time_heuristic(
     edep, e_units = units.unwrap_lgdo(edep)
 
     # we want to attach the right units to the dt heuristic, if possible
-    attrs = {}
+    unit = None
     if t_units is not None and e_units is not None:
-        attrs["units"] = units.unit_to_lh5_attr(t_units / e_units)
+        unit = units.unit_to_lh5_attr(t_units / e_units)
+        return units.attach_units(ak.Array(_drift_time_heuristic_impl(drift_time, edep)), unit)
 
-    return Array(_drift_time_heuristic_impl(drift_time, edep), attrs=attrs)
+    return ak.Array(_drift_time_heuristic_impl(drift_time, edep))
 
 
 @numba.njit(cache=True)
@@ -754,9 +753,9 @@ def _estimate_current_impl(
 
 
 def maximum_current(
-    edep: ArrayLike,
-    drift_time: ArrayLike,
-    dist_to_nplus: ArrayLike | None = None,
+    edep: ak.Array,
+    drift_time: ak.Array,
+    dist_to_nplus: ak.Array | None = None,
     *,
     template: np.array,
     times: np.array,
@@ -765,22 +764,22 @@ def maximum_current(
     activeness_surface: ArrayLike | None = None,
     surface_step_in_um: float = 10,
     return_mode: str = "current",
-) -> Array:
+) -> ak.Array:
     """Estimate the maximum current in the HPGe detector based on :func:`_estimate_current_impl`.
 
     Parameters
     ----------
     edep
-        Array of energies for each step.
+       Energies for each step.
     drift_time
-        Array of drift times for each step.
+        Drift times for each step.
     dist_to_nplus
         Distance to n-plus electrode, only needed if surface heuristics are enabled.
     template
-        array of the bulk pulse template
+        Array of the bulk pulse template
     times
         time-stamps for the bulk pulse template
-    fccd
+    fccd_in_um
         Value of the full-charge-collection depth, if `None` no surface corrections are performed.
     surface_library
         2D array (distance, time) of the rate of charge arriving at the p-n junction. Each row
@@ -797,9 +796,9 @@ def maximum_current(
     """
     # extract LGDO data and units
 
-    drift_time, _ = units.unwrap_lgdo(drift_time)
-    edep, _ = units.unwrap_lgdo(edep)
-    dist_to_nplus, _ = units.unwrap_lgdo(dist_to_nplus)
+    drift_time = units.units_conv_ak(drift_time, "ns")
+    edep = units.units_conv_ak(edep, "keV")
+    dist_to_nplus = units.units_conv_ak(dist_to_nplus, "um")
 
     include_surface_effects = False
 
@@ -837,11 +836,12 @@ def maximum_current(
 
     # return
     if return_mode == "max_time":
-        return Array(time)
+        return units.attach_units(ak.Array(time), "ns")
     if return_mode == "current":
-        return Array(curr)
+        # current has no unit (depends on the template)
+        return ak.Array(curr)
     if return_mode == "energy":
-        return Array(energy)
+        return units.attach_units(ak.Array(energy), "keV")
 
     msg = f"Return mode {return_mode} is not implemented."
     raise ValueError(msg)
diff --git a/src/reboost/hpge/surface.py b/src/reboost/hpge/surface.py
index 038f08e..bcb6739 100644
--- a/src/reboost/hpge/surface.py
+++ b/src/reboost/hpge/surface.py
@@ -5,27 +5,30 @@
 import awkward as ak
 import numba
 import numpy as np
-import pygeomhpges
-from lgdo import VectorOfVectors
+import pint
+import pyg4ometry
 from lgdo.types import LGDO
-from numpy.typing import ArrayLike
+from numpy.typing import NDArray
 from scipy import stats
 
+from reboost.units import ureg as u
+
+from .. import units
+
 log = logging.getLogger(__name__)
 
 
 def distance_to_surface(
-    positions_x: VectorOfVectors,
-    positions_y: VectorOfVectors,
-    positions_z: VectorOfVectors,
-    hpge: pygeomhpges.base.HPGe,
-    det_pos: ArrayLike,
+    positions_x: ak.Array,
+    positions_y: ak.Array,
+    positions_z: ak.Array,
+    hpge: legendhpges.base.HPGe,
+    det_pos: pint.Quantity | pyg4ometry.gdml.Position | tuple = (0, 0, 0) * u.m,
     *,
     surface_type: str | None = None,
-    unit: str = "mm",
-    distances_precompute: VectorOfVectors | None = None,
+    distances_precompute: ak.Array | None = None,
     precompute_cutoff: float | None = None,
-) -> VectorOfVectors:
+) -> ak.Array:
     """Computes the distance from each step to the detector surface.
 
     The calculation can be performed for any surface type `nplus`, `pplus`,
@@ -44,32 +47,36 @@ def distance_to_surface(
     hpge
         HPGe object.
     det_pos
-        position of the detector origin, must be a 3 component array corresponding to `(x,y,z)`.
+        position of the detector origin, must be a 3 component array corresponding to `(x,y,z)`. If no units
+        are specified mm is assumed.
     surface_type
         string of which surface to use, can be `nplus`, `pplus` `passive` or None (in which case the distance to any surface is calculated).
     unit
         unit for the hit tier positions table.
     distances_precompute
-        VectorOfVectors of distance to any surface computed by remage.
+        Distance to any surface computed by remage.
     precompute_cutoff
         cutoff on distances_precompute to not compute the distance for (in mm)
 
     Returns
     -------
-    VectorOfVectors with the same shape as `positions_x/y/z` of the distance to the surface.
+    Distance to the surface for each hit with the same shape as `positions_x/y/z`.
 
     Note
     ----
     `positions_x/positions_y/positions_z` must all have the same shape.
     """
-    factor = np.array([1, 100, 1000])[unit == np.array(["mm", "cm", "m"])][0]
-
     # compute local positions
     pos = []
     sizes = []
 
+    if not isinstance(det_pos, pint.Quantity | pyg4ometry.gdml.Position):
+        det_pos = det_pos * u.mm
+
+    det_pos = units.pg4_to_pint(det_pos)
+
     for idx, pos_tmp in enumerate([positions_x, positions_y, positions_z]):
-        local_pos_tmp = ak.Array(pos_tmp) * factor - det_pos[idx]
+        local_pos_tmp = units.units_conv_ak(pos_tmp, "mm") - det_pos[idx].to("mm").m
         local_pos_flat_tmp = ak.flatten(local_pos_tmp).to_numpy()
         pos.append(local_pos_flat_tmp)
         sizes.append(ak.num(local_pos_tmp, axis=1))
@@ -106,7 +113,7 @@ def distance_to_surface(
             local_positions[indices], surface_indices=surface_indices
         )
 
-    return VectorOfVectors(ak.unflatten(distances, size), dtype=np.float32)
+    return units.attach_units(ak.Array(ak.unflatten(distances, size)), "mm")
 
 
 @numba.njit(cache=True)
@@ -222,7 +229,7 @@ def get_surface_response(
     factor: float = 0.29,
     nsteps: int = 10000,
     delta_x: float = 10,
-):
+) -> NDArray:
     """Extract the surface response current pulse based on diffusion.
 
     This extracts the amount of charge arrived (cumulative) at the p-n
@@ -248,6 +255,10 @@ def get_surface_response(
         the number of time steps.
     delta_x
         the width of each position bin.
+
+    Returns
+    -------
+    2D array of the amount of charge arriving at the p-n junction as a function of time for each depth.
     """
     # number of position steps
     nx = int(fccd / delta_x)
diff --git a/src/reboost/math/functions.py b/src/reboost/math/functions.py
index 4368dfc..1a57e8d 100644
--- a/src/reboost/math/functions.py
+++ b/src/reboost/math/functions.py
@@ -5,14 +5,13 @@
 import awkward as ak
 import numpy as np
 from lgdo import Array, VectorOfVectors
-from lgdo.types import LGDO
+
+from .. import units
 
 log = logging.getLogger(__name__)
 
 
-def piecewise_linear_activeness(
-    distances: VectorOfVectors | ak.Array, fccd: float, dlf: float
-) -> VectorOfVectors | Array:
+def piecewise_linear_activeness(distances: ak.Array, fccd_in_mm: float, dlf: float) -> ak.Array:
     r"""Piecewise linear HPGe activeness model.
 
     Based on:
@@ -38,11 +37,10 @@ def piecewise_linear_activeness(
     Parameters
     ----------
     distances
-        the distance from each step to the detector surface. Can be either a
-        `numpy` or `awkward` array, or a LGDO `VectorOfVectors` or `Array`. The computation
+        the distance from each step to the detector surface. The computation
         is performed for each element and the shape preserved in the output.
 
-    fccd
+    fccd_in_mm
         the value of the FCCD
     dlf
         the fraction of the FCCD which is fully inactive.
@@ -52,14 +50,11 @@ def piecewise_linear_activeness(
     a :class:`VectorOfVectors` or :class:`Array` of the activeness
     """
     # convert to ak
-    if isinstance(distances, LGDO):
-        distances_ak = distances.view_as("ak")
-    elif not isinstance(distances, ak.Array):
-        distances_ak = ak.Array(distances)
-    else:
-        distances_ak = distances
+    distances = ak.Array(distances)
 
-    dl = fccd * dlf
+    distances_ak = units.units_conv_ak(distances, "mm")
+
+    dl = fccd_in_mm * dlf
     distances_flat = (
         ak.flatten(distances_ak).to_numpy() if distances_ak.ndim > 1 else distances_ak.to_numpy()
     )
@@ -68,24 +63,24 @@ def piecewise_linear_activeness(
     results = np.full_like(distances_flat, np.nan, dtype=np.float64)
     lengths = ak.num(distances_ak) if distances_ak.ndim > 1 else len(distances_ak)
 
-    mask1 = (distances_flat > fccd) | np.isnan(distances_flat)
+    mask1 = (distances_flat > fccd_in_mm) | np.isnan(distances_flat)
     mask2 = (distances_flat <= dl) & (~mask1)
     mask3 = ~(mask1 | mask2)
 
     # assign the values
     results[mask1] = 1
     results[mask2] = 0
-    results[mask3] = (distances_flat[mask3] - dl) / (fccd - dl)
+    results[mask3] = (distances_flat[mask3] - dl) / (fccd_in_mm - dl)
 
     # reshape
     results = ak.unflatten(ak.Array(results), lengths) if distances_ak.ndim > 1 else results
 
-    return VectorOfVectors(results) if results.ndim > 1 else Array(results)
+    return units.attach_units(results, "mm")
 
 
 def vectorised_active_energy(
-    distances: VectorOfVectors | ak.Array,
-    edep: VectorOfVectors | ak.Array,
+    distances: ak.Array,
+    edep: ak.Array,
     fccd: float | list,
     dlf: float | list,
 ) -> VectorOfVectors | Array:
@@ -115,7 +110,7 @@ def vectorised_active_energy(
 
     Returns
     -------
-    a :class:`VectorOfVectors` or :class:`Array` of the activeness
+    Activeness for each set of parameters
     """
     # add checks on fccd, dlf
     fccd = np.array(fccd)
@@ -133,20 +128,14 @@ def vectorised_active_energy(
 
     dl = fccd * dlf
 
-    def _convert(field):
+    def _convert(field, unit):
         # convert to ak
-        if isinstance(field, VectorOfVectors):
-            field_ak = field.view_as("ak")
-        elif not isinstance(field, ak.Array):
-            field_ak = ak.Array(field)
-        else:
-            msg = f"{field} must be an awkward array or VectorOfVectors"
-            raise TypeError(msg)
+        field_ak = units.units_conv_ak(field, unit)
 
         return field_ak, ak.flatten(field_ak).to_numpy()[:, np.newaxis]
 
-    distances_ak, distances_flat = _convert(distances)
-    _, edep_flat = _convert(edep)
+    distances_ak, distances_flat = _convert(distances, "mm")
+    _, edep_flat = _convert(edep, "keV")
     runs = ak.num(distances_ak, axis=-1)
 
     # vectorise fccd or tl
@@ -172,4 +161,4 @@ def _convert(field):
 
     energy = ak.sum(ak.unflatten(results * edep_flat, runs), axis=-2)
 
-    return VectorOfVectors(energy) if energy.ndim > 1 else Array(energy.to_numpy())
+    return units.attach_units(energy, "keV")
diff --git a/src/reboost/math/stats.py b/src/reboost/math/stats.py
index 0b78f53..6bd4f91 100644
--- a/src/reboost/math/stats.py
+++ b/src/reboost/math/stats.py
@@ -6,7 +6,6 @@
 import awkward as ak
 import numpy as np
 from lgdo import Array
-from numpy.typing import ArrayLike
 
 log = logging.getLogger(__name__)
 
@@ -72,7 +71,7 @@ def apply_energy_resolution(
     return ak.unflatten(energies_flat_smear, num)
 
 
-def gaussian_sample(mu: ArrayLike, sigma: ArrayLike | float, *, seed: int | None = None) -> Array:
+def gaussian_sample(mu: ak.Array, sigma: ak.Array | float, *, seed: int | None = None) -> ak.Array:
     r"""Generate samples from a gaussian.
 
     Based on:
@@ -99,9 +98,7 @@ def gaussian_sample(mu: ArrayLike, sigma: ArrayLike | float, *, seed: int | None
     """
     # convert inputs
 
-    if isinstance(mu, Array):
-        mu = mu.view_as("np")
-    elif isinstance(mu, ak.Array):
+    if isinstance(mu, ak.Array):
         mu = mu.to_numpy()
     elif not isinstance(mu, np.ndarray):
         mu = np.array(mu)
@@ -116,4 +113,4 @@ def gaussian_sample(mu: ArrayLike, sigma: ArrayLike | float, *, seed: int | None
 
     rng = np.random.default_rng(seed=seed)  # Create a random number generator
 
-    return Array(rng.normal(loc=mu, scale=sigma))
+    return ak.Array(rng.normal(loc=mu, scale=sigma))
diff --git a/src/reboost/shape/cluster.py b/src/reboost/shape/cluster.py
index 1445d85..f57b9d6 100644
--- a/src/reboost/shape/cluster.py
+++ b/src/reboost/shape/cluster.py
@@ -7,12 +7,12 @@
 import numpy as np
 from lgdo import VectorOfVectors
 
+from .. import units
+
 log = logging.getLogger(__name__)
 
 
-def apply_cluster(
-    cluster_run_lengths: VectorOfVectors | ak.Array, field: ak.Array | VectorOfVectors
-) -> VectorOfVectors:
+def apply_cluster(cluster_run_lengths: ak.Array, field: ak.Array) -> ak.Array:
     """Apply clustering to a field.
 
     Parameters
@@ -32,18 +32,18 @@ def apply_cluster(
     clusters = ak.unflatten(ak.flatten(field), ak.flatten(cluster_run_lengths))
 
     # reshape into cluster oriented
-    return VectorOfVectors(ak.unflatten(clusters, n_cluster))
+    return ak.Array(ak.unflatten(clusters, n_cluster), attrs=field.attrs)
 
 
 def cluster_by_step_length(
-    trackid: ak.Array | VectorOfVectors,
-    pos_x: ak.Array | VectorOfVectors,
-    pos_y: ak.Array | VectorOfVectors,
-    pos_z: ak.Array | VectorOfVectors,
-    dist: ak.Array | VectorOfVectors | None = None,
+    trackid: ak.Array,
+    pos_x: ak.Array,
+    pos_y: ak.Array,
+    pos_z: ak.Array,
+    dist: ak.Array | None = None,
     surf_cut: float | None = None,
-    threshold: float = 0.1,
-    threshold_surf: float | None = None,
+    threshold_in_mm: float = 0.1,
+    threshold_surf_in_mm: float | None = None,
 ) -> VectorOfVectors:
     """Perform clustering based on the step length.
 
@@ -59,20 +59,20 @@ def cluster_by_step_length(
     Parameters
     ----------
     trackid
-        index of the track.
+        index of the tracks.
     pos_x
-        x position of the step.
+        x position of the steps.
     pos_y
-        y position of the step.
+        y position of the steps.
     pos_z
-        z position of the step.
+        z position of the steps.
     dist
         distance to the detector surface. Can be `None` in which case all steps are treated as being in the "bulk".
     surf_cut
         Size of the surface region (in mm), if `None` no selection is applied (default).
-    threshold
+    threshold_in_mm
         Distance threshold in mm to combine steps in the bulk.
-    threshold_surf
+    threshold_surf_in_mm
         Distance threshold in mm to combine steps in the surface.
 
     Returns
@@ -80,37 +80,22 @@ def cluster_by_step_length(
     Array of the run lengths of each cluster within a hit.
     """
     # type conversions
-    if isinstance(pos_x, VectorOfVectors):
-        pos_x = pos_x.view_as("ak")
-
-    if isinstance(pos_y, VectorOfVectors):
-        pos_y = pos_y.view_as("ak")
-
-    if isinstance(pos_z, VectorOfVectors):
-        pos_z = pos_z.view_as("ak")
-
-    if isinstance(trackid, VectorOfVectors):
-        trackid = trackid.view_as("ak")
-
-    if isinstance(dist, VectorOfVectors):
-        dist = dist.view_as("ak")
 
     pos = np.vstack(
         [
-            ak.flatten(pos_x).to_numpy().astype(np.float64),
-            ak.flatten(pos_y).to_numpy().astype(np.float64),
-            ak.flatten(pos_z).to_numpy().astype(np.float64),
+            ak.flatten(units.units_conv_ak(p, "mm")).to_numpy().astype(np.float64)
+            for p in [pos_x, pos_y, pos_z]
         ]
     ).T
 
-    indices_flat = cluster_by_distance_numba(
+    indices_flat = _cluster_by_distance_numba(
         ak.flatten(ak.local_index(trackid)).to_numpy(),
         ak.flatten(trackid).to_numpy(),
         pos,
         dist_to_surf=ak.flatten(dist).to_numpy() if dist is not None else dist,
         surf_cut=surf_cut,
-        threshold=threshold,
-        threshold_surf=threshold_surf,
+        threshold=threshold_in_mm,
+        threshold_surf=threshold_surf_in_mm,
     )
 
     # reshape into being event oriented
@@ -119,11 +104,11 @@ def cluster_by_step_length(
     # number of steps per cluster
     counts = ak.run_lengths(indices)
 
-    return VectorOfVectors(counts)
+    return ak.Array(counts)
 
 
 @numba.njit
-def cluster_by_distance_numba(
+def _cluster_by_distance_numba(
     local_index: np.ndarray,
     trackid: np.ndarray,
     pos: np.ndarray,
@@ -210,10 +195,10 @@ def _dist(a, b):
 
 
 def step_lengths(
-    x_cluster: ak.Array | VectorOfVectors,
-    y_cluster: ak.Array | VectorOfVectors,
-    z_cluster: ak.Array | VectorOfVectors,
-) -> VectorOfVectors:
+    x_cluster: ak.Array,
+    y_cluster: ak.Array,
+    z_cluster: ak.Array,
+) -> ak.Array:
     """Compute the distance between consecutive steps.
 
     This is based on calculating the distance between consecutive steps in the same track,
@@ -242,13 +227,14 @@ def step_lengths(
     data = [x_cluster, y_cluster, z_cluster]
 
     for idx, var in enumerate(data):
-        if isinstance(var, VectorOfVectors):
-            data[idx] = var.view_as("ak")
         # check shape
-        if data[idx].ndim != 3:
+        if var.ndim != 3:
             msg = f"The input array for step lengths must be 3 dimensional not {data[idx.dim]}"
             raise ValueError(msg)
 
+        # type convert
+        data[idx] = units.units_conv_ak(data[idx], "mm")
+
     counts = ak.num(data[0], axis=-1)
     data = np.vstack([ak.flatten(ak.flatten(var)).to_numpy() for var in data])
     dist = np.append(np.sqrt(np.sum(np.diff(data, axis=1) ** 2, axis=0)), 0)
@@ -257,4 +243,4 @@ def step_lengths(
     clusters = ak.unflatten(ak.Array(dist), ak.flatten(counts))
 
     out = ak.unflatten(clusters, n_cluster)
-    return VectorOfVectors(out[:, :, :-1])
+    return ak.Array(out[:, :, :-1], attrs={"units": "mm"})
diff --git a/src/reboost/units.py b/src/reboost/units.py
index 1446283..2d9e466 100644
--- a/src/reboost/units.py
+++ b/src/reboost/units.py
@@ -45,6 +45,19 @@ def units_convfact(data: Any | LGDO | ak.Array, target_units: pint.Unit | str) -
     return 1
 
 
+def attach_units(data: ak.Array, unit: str) -> ak.Array:
+    """Convenience function to attach units to `ak.Array`.
+
+    Parameters
+    ----------
+    data
+        the array to add units to
+    unit
+        the unit
+    """
+    return ak.with_parameter(data, parameter="units", value=unit)
+
+
 def units_conv_ak(data: Any | LGDO | ak.Array, target_units: pint.Unit | str) -> ak.Array:
     """Calculate numeric conversion factor to reach `target_units`, and apply to data converted to ak.
 
@@ -99,6 +112,25 @@ def unwrap_lgdo(data: Any | LGDO | ak.Array, library: str = "ak") -> tuple[Any,
     return ret_data, ret_units
 
 
+def get_unit_str(data: ak.Array) -> str:
+    """Get the units as a string for an awkward array with attached units.
+
+    Parameters
+    ----------
+    data
+        the array with attached units
+
+    Returns
+    -------
+    a string of the units.
+    """
+    attrs = ak.parameters(data)
+
+    if "units" in attrs:
+        return attrs["units"]
+    return None
+
+
 def unit_to_lh5_attr(unit: pint.Unit) -> str:
     """Convert Pint unit to a string that can be used as attrs["units"] in an LGDO."""
     # TODO: we should check if this can be always parsed by Unitful.jl
diff --git a/tests/hit/configs/hit_config.yaml b/tests/hit/configs/hit_config.yaml
index 98945c4..c13b8f4 100644
--- a/tests/hit/configs/hit_config.yaml
+++ b/tests/hit/configs/hit_config.yaml
@@ -36,7 +36,7 @@ processing_groups:
 
       activeness:
         reboost.math.functions.piecewise_linear_activeness(HITS.distance_to_nplus,
-        fccd=DETECTOR_OBJECTS.det_pars.fccd_in_mm,
+        fccd_in_mm=DETECTOR_OBJECTS.det_pars.fccd_in_mm,
         dlf=DETECTOR_OBJECTS.det_pars.dlf)
 
       active_energy: ak.sum(HITS.edep*HITS.activeness, axis=-1)
diff --git a/tests/hpge/test_current.py b/tests/hpge/test_current.py
index a877b6d..9870f8b 100644
--- a/tests/hpge/test_current.py
+++ b/tests/hpge/test_current.py
@@ -3,8 +3,8 @@
 import awkward as ak
 import numpy as np
 import pytest
-from lgdo import Array, VectorOfVectors
 
+from reboost import units
 from reboost.hpge import psd, surface
 from reboost.shape import cluster
 
@@ -49,15 +49,11 @@ def test_model():
 def test_maximum_current(test_model):
     model, x = test_model
 
-    edep = VectorOfVectors(
-        ak.Array([[100.0, 300.0, 50.0], [10.0, 0.0, 100.0], [500.0]]), attrs={"unit": "keV"}
-    )
-    times = VectorOfVectors(
-        ak.Array([[400, 500, 700], [800, 0, 1500], [700]], attrs={"unit": "ns"})
-    )
+    edep = units.attach_units(ak.Array([[100.0, 300.0, 50.0], [10.0, 0.0, 100.0], [500.0]]), "keV")
+    times = units.attach_units(ak.Array([[400, 500, 700], [800, 0, 1500], [700]]), "ns")
 
     curr = psd.maximum_current(edep, times, template=model, times=x)
-    assert isinstance(curr, Array)
+    assert isinstance(curr, ak.Array)
 
     assert len(curr) == 3
 
@@ -73,7 +69,7 @@ def test_maximum_current(test_model):
         return_mode="max_time",
     )
 
-    assert isinstance(max_t, Array)
+    assert isinstance(max_t, ak.Array)
     assert len(max_t) == 3
 
     # should be close to 250 (could be some differences due to the discretisation)
@@ -87,41 +83,68 @@ def test_maximum_current(test_model):
         return_mode="energy",
     )
 
-    assert isinstance(energy, Array)
+    assert isinstance(energy, ak.Array)
     assert len(energy) == 3
 
     # should be close to 250 (could be some differences due to the discretisation)
     assert abs(energy[2] - 500.0) < 2
 
 
-def test_with_cluster(test_model):
+def test_units(test_model):
+    # standard units
     model, x = test_model
 
-    edep = VectorOfVectors(
-        ak.Array([[100.0, 300.0, 50.0], [10.0, 1.0, 100.0], [500.0]]), attrs={"unit": "keV"}
+    edep = units.attach_units(ak.Array([[100.0, 300.0, 50.0], [10.0, 0.0, 100.0], [500.0]]), "keV")
+    times = units.attach_units(ak.Array([[400, 500, 700], [800, 0, 1500], [700]]), "ns")
+
+    energy = psd.maximum_current(edep, times, template=model, times=x, return_mode="energy")
+    unit = units.get_unit_str(energy)
+
+    assert unit == "keV"
+
+    # try for time
+    time = psd.maximum_current(edep, times, template=model, times=x, return_mode="max_time")
+
+    unit = units.get_unit_str(time)
+    assert unit == "ns"
+
+    # try with different units and check the result is the same
+    times_us = units.attach_units(
+        ak.Array([[0.400, 0.500, 0.700], [0.800, 0, 1.500], [0.700]]), "us"
     )
-    times = VectorOfVectors(
-        ak.Array([[400, 410, 420], [800, 0, 1500], [700]], attrs={"unit": "ns"})
+    max_time_us = psd.maximum_current(
+        edep, times_us, template=model, times=x, return_mode="max_time"
     )
-    xloc = VectorOfVectors(ak.Array([[1, 1.1, 1.2], [0, 50, 80], [100]], attrs={"unit": "mm"}))
 
-    dist = VectorOfVectors(ak.Array([[50, 40, 0.2], [300, 0.4, 0.2], [0.8]], attrs={"unit": "ns"}))
+    unit = units.get_unit_str(max_time_us)
+    assert unit == "ns"
+
+    assert ak.all(time == max_time_us)
+
+
+def test_with_cluster(test_model):
+    model, x = test_model
+
+    edep = units.attach_units(ak.Array([[100.0, 300.0, 50.0], [10.0, 1.0, 100.0], [500.0]]), "keV")
+    times = units.attach_units(ak.Array([[400, 410, 420], [800, 0, 1500], [700]]), "ns")
+    xloc = units.attach_units(ak.Array([[1, 1.1, 1.2], [0, 50, 80], [100]]), "mm")
+    dist = units.attach_units(ak.Array([[50, 40, 0.2], [300, 0.4, 0.2], [0.8]]), "mm")
 
     yloc = ak.full_like(xloc, 0.0)
     zloc = ak.full_like(xloc, 0.0)
     trackid = ak.full_like(xloc, 0)
 
     clusters = cluster.cluster_by_step_length(
-        trackid, xloc, yloc, zloc, dist, threshold=1, threshold_surf=1, surf_cut=0
+        trackid, xloc, yloc, zloc, dist, threshold_in_mm=1, threshold_surf_in_mm=1, surf_cut=0
     )
-    cluster_edep = cluster.apply_cluster(clusters, edep).view_as("ak")
-    cluster_times = cluster.apply_cluster(clusters, times).view_as("ak")
+    cluster_edep = cluster.apply_cluster(clusters, edep)
+    cluster_times = cluster.apply_cluster(clusters, times)
 
     e = ak.sum(cluster_edep, axis=-1)
     t = ak.sum(cluster_edep * cluster_times, axis=-1) / e
     curr = psd.maximum_current(e, t, template=model, times=x)
 
-    assert isinstance(curr, Array)
+    assert isinstance(curr, ak.Array)
     assert len(curr) == 3
 
     # should be close to 250 (could be some differences due to the discretisation)
@@ -133,21 +156,23 @@ def test_maximum_current_surface(test_model):
 
     # test for both input types
     for dtype in [np.float64, np.float32]:
-        edep = VectorOfVectors(
-            ak.values_astype(ak.Array([[100.0, 300.0, 50.0], [10.0, 0.0, 100.0], [500.0]]), dtype),
-            attrs={"unit": "keV"},
+        edep = units.attach_units(
+            ak.Array(
+                ak.values_astype(
+                    ak.Array([[100.0, 300.0, 50.0], [10.0, 0.0, 100.0], [500.0]]), dtype
+                ),
+            ),
+            "keV",
         )
 
-        times = VectorOfVectors(
-            ak.values_astype(
-                ak.Array([[400, 500, 700], [800, 0, 1500], [700]], attrs={"unit": "ns"}), dtype
-            )
+        times = units.attach_units(
+            ak.Array(ak.values_astype(ak.Array([[400, 500, 700], [800, 0, 1500], [700]]), dtype)),
+            "ns",
         )
 
-        dist = VectorOfVectors(
-            ak.values_astype(
-                ak.Array([[50, 40, 0.2], [300, 0.4, 0.2], [0.8]], attrs={"unit": "ns"}), dtype
-            )
+        dist = units.attach_units(
+            ak.Array(ak.values_astype(ak.Array([[50, 40, 0.2], [300, 0.4, 0.2], [0.8]]), dtype)),
+            "mm",
         )
 
         surface_models = surface.get_surface_library(1002, 10)
@@ -167,7 +192,7 @@ def test_maximum_current_surface(test_model):
             activeness_surface=surface_activeness,
             times=x,
             return_mode="current",
-        ).view_as("np")
+        )
 
         curr_bulk = psd.maximum_current(
             edep,
@@ -176,7 +201,7 @@ def test_maximum_current_surface(test_model):
             template=model,
             times=x,
             return_mode="current",
-        ).view_as("np")
+        )
         # check shape
 
         assert len(curr_surf) == 3
diff --git a/tests/hpge/test_dt_heuristic.py b/tests/hpge/test_dt_heuristic.py
index f70bc44..e735c53 100644
--- a/tests/hpge/test_dt_heuristic.py
+++ b/tests/hpge/test_dt_heuristic.py
@@ -1,19 +1,20 @@
 from __future__ import annotations
 
+import awkward as ak
 import numpy as np
 import pytest
 from dbetto import AttrsDict
-from lgdo import Array, Table, VectorOfVectors, lh5
+from lgdo import lh5
 from scipy.interpolate import RegularGridInterpolator
 
+from reboost import units
 from reboost.hpge import psd, utils
 from reboost.units import ureg as u
 
 # data from remage sim in ./simulation
 # fmt: off
-gamma_stp = Table(
-    {
-        "edep": VectorOfVectors(
+gamma_stp = ak.Array({
+        "edep": units.attach_units(ak.Array(
             [
                 [763, 20.1, 30.2, 40.5, 36.5, 110],
                 [0.0158, 526, 111],
@@ -26,9 +27,8 @@
                 [728, 16.1, 256],
                 [613, 0.179, 9.03, 101, 26.9, 22.9, 26.1, 31.5, 11.1, 5.56, 31.7, 27.2, 23, 33.7, 37.6],
             ],
-            attrs={"units": "keV"},
-        ),
-        "xloc": VectorOfVectors(
+        ), "keV"),
+        "xloc": units.attach_units(ak.Array(
             [
                 [-0.0022, 0.0115, 0.0141, 0.0169, 0.0165, 0.0165],
                 [0.0106, 0.0104, 0.0103],
@@ -41,9 +41,9 @@
                 [-0.00369, -0.00343, 0.00682],
                 [-0.0162, -0.0224, -0.0224, -0.0224, -0.0224, -0.0224, -0.0224, -0.0224, -0.022, -0.022, -0.022, -0.022, -0.022, -0.022, -0.022],
             ],
-            attrs={"units": "m"},
-        ),
-        "yloc": VectorOfVectors(
+
+        ), "m"),
+        "yloc": units.attach_units(ak.Array(
             [
                 [-0.00951, -0.0201, -0.0199, -0.014, -0.00924, -0.0092],
                 [-0.0264, -0.0264, -0.0264],
@@ -56,9 +56,8 @@
                 [-0.000536, -0.000655, -0.00369],
                 [-0.00355, -0.0134, -0.0134, -0.0134, -0.0134, -0.0134, -0.0134, -0.0134, -0.0133, -0.0133, -0.0133, -0.0133, -0.0133, -0.0133, -0.0133],
             ],
-            attrs={"units": "m"},
-        ),
-        "zloc": VectorOfVectors(
+        ),"m"),
+        "zloc": units.attach_units(ak.Array(
             [
                 [0.0281, 0.0124, 0.0112, 0.00313, 0.0052, 0.00495],
                 [0.029, 0.0292, 0.0293],
@@ -71,8 +70,7 @@
                 [0.0197, 0.0195, 0.0117],
                 [0.0029, 0.00108, 0.00108, 0.00107, 0.00106, 0.00106, 0.00106, 0.00106, 0.00128, 0.00128, 0.00129, 0.0013, 0.0013, 0.0013, 0.0013],
             ],
-            attrs={"units": "m"},
-        ),
+        ), "m"),
     }
 )
 # fmt: on
@@ -123,12 +121,12 @@ def dt_map_dummy(legendtestdata):
 
 
 def test_drift_time_dummy(dt_map_dummy):
-    gamma_stp_shift = Table(
+    gamma_stp_shift = ak.Array(
         {
             "edep": gamma_stp.edep,
-            "xloc": VectorOfVectors(gamma_stp.xloc.view_as("ak") + 10, attrs={"units": "m"}),
-            "yloc": VectorOfVectors(gamma_stp.yloc.view_as("ak") + 10, attrs={"units": "m"}),
-            "zloc": VectorOfVectors(gamma_stp.zloc.view_as("ak") + 10, attrs={"units": "m"}),
+            "xloc": ak.Array(gamma_stp.xloc + 10),
+            "yloc": ak.Array(gamma_stp.yloc + 10),
+            "zloc": ak.Array(gamma_stp.zloc + 10),
         }
     )
 
@@ -146,11 +144,11 @@ def test_drift_time_dummy(dt_map_dummy):
         * u.m,
     )
     # turn into an Awkward array so we can index
-    dt_arr = dt_values.view_as("ak")
+    dt_arr = dt_values
 
     # helper to pull out expected from the RegularGridInterpolator
     def expected_dt(event, step):
-        stp = gamma_stp.view_as("ak")
+        stp = gamma_stp
         x = stp.xloc[event][step]
         y = stp.yloc[event][step]
         z = stp.zloc[event][step]
@@ -176,20 +174,12 @@ def test_drift_time(dt_map):
         gamma_stp.zloc,
         dt_map,
     )
-    assert isinstance(dt_values, VectorOfVectors)
+
+    dt_values, unit = units.unwrap_lgdo(dt_values)
+
+    assert isinstance(dt_values, ak.Array)
     assert dt_values.ndim == 2
-    assert dt_values.attrs["units"] == "ns"
-
-    # test whether this works with non-LGDOs
-    data = gamma_stp.view_as("ak")
-    dt_values_nolgdo = psd.drift_time(
-        data.xloc,  # units should match the dt map units -> meters
-        data.yloc,
-        data.zloc,
-        dt_map,
-    )
-    assert isinstance(dt_values_nolgdo, VectorOfVectors)
-    assert dt_values_nolgdo == dt_values
+    assert units.unit_to_lh5_attr(unit) == "ns"
 
 
 def test_drift_time_heuristics(dt_map):
@@ -204,6 +194,7 @@ def test_drift_time_heuristics(dt_map):
         dt_values,
         gamma_stp.edep,
     )
+    dt_heu, unit = units.unwrap_lgdo(dt_heu)
+    assert isinstance(dt_heu, ak.Array)
 
-    assert isinstance(dt_heu, Array)
-    assert dt_heu.attrs["units"] == "ns/keV"
+    assert units.unit_to_lh5_attr(unit) == "ns/keV"
diff --git a/tests/hpge/test_r90.py b/tests/hpge/test_r90.py
index f184b82..e1b9974 100644
--- a/tests/hpge/test_r90.py
+++ b/tests/hpge/test_r90.py
@@ -116,8 +116,8 @@ def test_r90():
     elect_r90 = 0.0001816
     gamma_r90 = 0.0157370
 
-    assert round(r90_output.nda[0], 7) == elect_r90
-    assert round(r90_output.nda[1], 7) == gamma_r90
+    assert round(r90_output[0], 7) == elect_r90
+    assert round(r90_output[1], 7) == gamma_r90
 
     edep = [gamma_edep]
     xloc = [gamma_xloc]
@@ -134,4 +134,4 @@ def test_r90():
     )
 
     r90_output = r90(data.edep, data.xloc, data.yloc, data.zloc)
-    assert round(r90_output.nda[0], 7) == gamma_r90
+    assert round(r90_output[0], 7) == gamma_r90
diff --git a/tests/hpge/test_surface.py b/tests/hpge/test_surface.py
index e05e535..5fe8210 100644
--- a/tests/hpge/test_surface.py
+++ b/tests/hpge/test_surface.py
@@ -5,11 +5,10 @@
 import pyg4ometry
 import pytest
 from legendtestdata import LegendTestData
-from lgdo import types
-from lgdo.types import VectorOfVectors
-from pygeomhpges import make_hpge
 
+from reboost import units
 from reboost.hpge.surface import distance_to_surface, get_surface_response
+from reboost.units import ureg as u
 
 
 @pytest.fixture(scope="session")
@@ -21,7 +20,7 @@ def test_data_configs():
 
 def test_distance_to_surface(test_data_configs):
     gedet = make_hpge(test_data_configs + "/V99000A.json", registry=pyg4ometry.geant4.Registry())
-    dist = [100, 0, 0]
+    dist = [100, 0, 0] * u.mm
 
     pos = ak.Array(
         {
@@ -47,21 +46,21 @@ def test_distance_to_surface(test_data_configs):
     dist_full = distance_to_surface(
         pos.xloc, pos.yloc, pos.zloc, gedet, det_pos=dist, surface_type=None
     )
-    assert isinstance(dist_full, types.LGDO)
+    assert isinstance(dist_full, ak.Array)
 
     # check skipping the calculation for points > 5 mm
     dist = distance_to_surface(
-        VectorOfVectors(pos.xloc),
-        VectorOfVectors(pos.yloc),
-        VectorOfVectors(pos.zloc),
+        pos.xloc,
+        pos.yloc,
+        pos.zloc,
         gedet,
         det_pos=dist,
         surface_type=None,
-        distances_precompute=VectorOfVectors(pos.distance_to_surface),
+        distances_precompute=pos.distance_to_surface,
         precompute_cutoff=5,
     )
 
-    assert isinstance(dist, types.LGDO)
+    assert isinstance(dist, ak.Array)
 
     assert ak.all(dist[0] == dist_full[0])
     assert ak.all(dist[2] == dist_full[2])
@@ -69,6 +68,42 @@ def test_distance_to_surface(test_data_configs):
     assert np.isnan(dist[1][0])
 
 
+def test_units(test_data_configs):
+    gedet = make_hpge(test_data_configs + "/V99000A.json", registry=pyg4ometry.geant4.Registry())
+    dist = [100, 0, 0] * u.mm
+
+    pos = ak.Array(
+        {
+            "xloc": units.attach_units([[0, 100, 200], [100], [700, 500, 200]], "mm"),
+            "yloc": units.attach_units([[100, 0, 0], [200], [100, 300, 200]], "mm"),
+            "zloc": units.attach_units([[700, 10, 20], [100], [300, 100, 0]], "mm"),
+            "distance_to_surface": units.attach_units([[1, 1, 1], [10], [1, 1, 1]], "m"),
+        }
+    )
+
+    dists = distance_to_surface(
+        pos.xloc, pos.yloc, pos.zloc, gedet, det_pos=dist, surface_type=None
+    )
+
+    assert units.get_unit_str(dists) == "mm"
+
+    pos_m = ak.Array(
+        {
+            "xloc": units.attach_units([[0, 0.100, 0.200], [0.100], [0.700, 0.500, 0.200]], "m"),
+            "yloc": units.attach_units([[0.100, 0, 0], [0.200], [0.100, 0.300, 0.200]], "m"),
+            "zloc": units.attach_units([[0.700, 0.01, 0.02], [0.100], [0.300, 0.100, 0]], "m"),
+            "distance_to_surface": units.attach_units([[1, 1, 1], [10], [1, 1, 1]], "m"),
+        }
+    )
+    dist_m = distance_to_surface(
+        pos_m.xloc, pos_m.yloc, pos_m.zloc, gedet, det_pos=dist, surface_type=None
+    )
+
+    assert units.get_unit_str(dist_m) == "mm"
+
+    assert ak.all(dists == dist_m)
+
+
 def test_surface_response():
     response = get_surface_response(fccd=1000, init=500)
 
diff --git a/tests/test_core.py b/tests/test_core.py
index f480b0c..13ccb92 100644
--- a/tests/test_core.py
+++ b/tests/test_core.py
@@ -180,7 +180,7 @@ def test_eval_expression():
     )
 
     # test with a reboost function
-    expression = "reboost.math.functions.piecewise_linear_activeness(HITS.distances, fccd=args.fccd, dlf=args.dlf)"
+    expression = "reboost.math.functions.piecewise_linear_activeness(HITS.distances, fccd_in_mm=args.fccd, dlf=args.dlf)"
 
     local = {"args": AttrsDict({"fccd": 1, "dlf": 0.5})}
 
diff --git a/tests/test_math.py b/tests/test_math.py
index fe2cb3c..6f1a4b2 100644
--- a/tests/test_math.py
+++ b/tests/test_math.py
@@ -12,7 +12,7 @@ def test_hpge_activeness():
     # test with VectorOfVectors and ak.Array input
     distances = ak.Array([[0.2], [0.6], [2]])
 
-    activeness = functions.piecewise_linear_activeness(VectorOfVectors(distances), fccd=1, dlf=0.4)
+    activeness = functions.piecewise_linear_activeness(ak.Array(distances), fccd_in_mm=1, dlf=0.4)
 
     # first point should be 0
     assert activeness[0][0] == 0
@@ -22,7 +22,7 @@ def test_hpge_activeness():
 
     # test with ak.Array input
     distances = ak.Array([[0.2], [0.6], [2]])
-    activeness = functions.piecewise_linear_activeness(distances, fccd=1, dlf=0.4)
+    activeness = functions.piecewise_linear_activeness(distances, fccd_in_mm=1, dlf=0.4)
 
     # first point should be 0
     assert activeness[0][0] == 0
@@ -31,19 +31,9 @@ def test_hpge_activeness():
     assert activeness[2][0] == 1
 
     # test with Array
-    activeness = functions.piecewise_linear_activeness([[0.2, 0.6, 2]], fccd=1, dlf=0.4)
+    activeness = functions.piecewise_linear_activeness([[0.2, 0.6, 2]], fccd_in_mm=1, dlf=0.4)
 
-    assert np.allclose(activeness.view_as("np"), [0, 1 / 3.0, 1])
-
-    # test with array
-    distances = [0.2, 0.6, 2]
-    activeness = functions.piecewise_linear_activeness(Array(distances), fccd=1, dlf=0.4)
-
-    # first point should be 0
-    assert activeness[0] == 0
-    # second should be 0.1/0.5 = 0.2
-    assert activeness[1] == pytest.approx(1 / 3.0)
-    assert activeness[2] == 1
+    assert np.allclose(activeness, [0, 1 / 3.0, 1])
 
 
 def test_vectorised_activeness():
@@ -52,18 +42,14 @@ def test_vectorised_activeness():
     edep = ak.Array([[100, 200], [100], [100], [100]])
 
     # simple case
-    energy = functions.vectorised_active_energy(
-        VectorOfVectors(distances), VectorOfVectors(edep), fccd=1, dlf=0.4
-    )
+    energy = functions.vectorised_active_energy(distances, edep, fccd=1, dlf=0.4)
 
     assert energy[0] == 200
     assert energy[1] == pytest.approx(100 / 3.0)
     assert energy[2] == 100
 
     # now vectorised over fccd
-    energy_fccd = functions.vectorised_active_energy(
-        VectorOfVectors(distances), VectorOfVectors(edep), fccd=[0, 0.5, 1, 2, 3], dlf=1
-    )
+    energy_fccd = functions.vectorised_active_energy(distances, edep, fccd=[0, 0.5, 1, 2, 3], dlf=1)
 
     # with distance of 0.2 only fccd of 0 gives energy
     assert np.allclose(energy_fccd[0], [300, 200, 200, 200, 200])
@@ -85,19 +71,19 @@ def test_vectorised_activeness():
 def test_sample():
     # list inputs
     samples = stats.gaussian_sample([1, 2, 3], [0.1, 0.1, 0.1])
-    assert isinstance(samples, Array)
+    assert isinstance(samples, ak.Array)
 
     # LGDO inputs
     samples = stats.gaussian_sample(Array(np.array([1, 2, 3])), Array(np.array([0.1, 0.1, 0.1])))
-    assert isinstance(samples, Array)
+    assert isinstance(samples, ak.Array)
 
     # ak inputs
     samples = stats.gaussian_sample(ak.Array([1, 2, 3]), ak.Array([1, 2, 3]))
-    assert isinstance(samples, Array)
+    assert isinstance(samples, ak.Array)
 
     # sigma float
     samples = stats.gaussian_sample([1, 2, 3], 0.1)
-    assert isinstance(samples, Array)
+    assert isinstance(samples, ak.Array)
 
 
 def test_energy_res():
diff --git a/tests/test_shape.py b/tests/test_shape.py
index 1b44136..3776780 100644
--- a/tests/test_shape.py
+++ b/tests/test_shape.py
@@ -118,7 +118,7 @@ def test_cluster_basic():
 
     runs = ak.run_lengths(trackid)
     assert ak.all(
-        cluster.apply_cluster(runs, trackid).view_as("ak")
+        cluster.apply_cluster(runs, trackid)
         == ak.Array([[[1, 1, 1], [2, 2], [3, 3], [7]], [[2, 2, 2], [3, 3, 3]], [[1]]])
     )
 
@@ -132,15 +132,15 @@ def test_cluster_by_step_length():
     dist = ak.Array([[0.1, 0.1, 0.1, 0.1, 2, 2, 2, 2], [2, 2, 2, 2, 2, 2], [0.3]])
 
     clusters = cluster.cluster_by_step_length(
-        trackid, x, y, z, dist, threshold=0.2, threshold_surf=0.02, surf_cut=0.15
+        trackid, x, y, z, dist, threshold_in_mm=0.2, threshold_surf_in_mm=0.02, surf_cut=0.15
     )
 
-    assert ak.all(clusters.view_as("ak") == ak.Array([[1, 1, 1, 1, 3, 1], [1, 1, 1, 2, 1], [1]]))
+    assert ak.all(clusters == ak.Array([[1, 1, 1, 1, 3, 1], [1, 1, 1, 2, 1], [1]]))
 
     # test without surface region
 
-    clusters = cluster.cluster_by_step_length(trackid, x, y, z, threshold=0.2)
-    assert ak.all(clusters.view_as("ak") == ak.Array([[2, 1, 1, 3, 1], [1, 1, 1, 2, 1], [1]]))
+    clusters = cluster.cluster_by_step_length(trackid, x, y, z, threshold_in_mm=0.2)
+    assert ak.all(clusters == ak.Array([[2, 1, 1, 3, 1], [1, 1, 1, 2, 1], [1]]))
 
 
 def test_step_length():
@@ -151,11 +151,11 @@ def test_step_length():
 
     run = ak.run_lengths(trackid)
 
-    x_cluster = cluster.apply_cluster(run, x).view_as("ak")
-    y_cluster = cluster.apply_cluster(run, y).view_as("ak")
-    z_cluster = cluster.apply_cluster(run, z).view_as("ak")
+    x_cluster = cluster.apply_cluster(run, x)
+    y_cluster = cluster.apply_cluster(run, y)
+    z_cluster = cluster.apply_cluster(run, z)
 
-    steps = cluster.step_lengths(x_cluster, y_cluster, z_cluster).view_as("ak")
+    steps = cluster.step_lengths(x_cluster, y_cluster, z_cluster)
 
     for idx, sub_array in enumerate(steps):
         assert ak.all(sub_array == ak.Array([[[0, 1], [1], [1], []], [[1, 3], [0, 1.0]], []])[idx])
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 45616f4..67d9d74 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -67,7 +67,9 @@ def test_get_function_string():
 
     # try a reboost package
 
-    expression = "reboost.math.functions.piecewise_linear_activeness(distances,fccd=1,dlf=0.5)"
+    expression = (
+        "reboost.math.functions.piecewise_linear_activeness(distances,fccd_in_mm=1,dlf=0.5)"
+    )
     func_string, globals_dict = get_function_string(expression)
 
     assert func_string == expression
@@ -76,11 +78,9 @@ def test_get_function_string():
     distances = awkward.Array([[0.2], [2], [0.6]])
 
     # just check it runs
-    assert eval(func_string, {"distances": distances}, globals_dict).view_as("ak")[0][0] == 0
-    assert eval(func_string, {"distances": distances}, globals_dict).view_as("ak")[1][0] == 1
-    assert eval(func_string, {"distances": distances}, globals_dict).view_as("ak")[2][
-        0
-    ] == pytest.approx(0.2)
+    assert eval(func_string, {"distances": distances}, globals_dict)[0][0] == 0
+    assert eval(func_string, {"distances": distances}, globals_dict)[1][0] == 1
+    assert eval(func_string, {"distances": distances}, globals_dict)[2][0] == pytest.approx(0.2)
 
     # try a more compliated expression
     expression = (