diff --git a/pyproject.toml b/pyproject.toml
index 1f63fa3..a4228c0 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -32,7 +32,7 @@ dependencies = ["numpy>=1.21.2", "scipy"]
 dynamic = ["version", "description"]
 
 [project.optional-dependencies]
-test = ["h5py", "tables < 3.8; python_version < '3.9'", "tables", "pandas", "uproot", "pytest-cov"]
+test = ["h5py", "tables < 3.8; python_version < '3.9'", "tables <= 3.9.1", "pandas", "uproot", "pytest-cov"]
 docs = ["sphinx","sphinx-rtd-theme","pandas"]
 dev = ["pytest","pre-commit","reuse","black","ruff","pylint","mypy"]
 
@@ -74,7 +74,7 @@ enable_error_code = ["ignore-without-code", "redundant-expr", "truthy-bool"]
 max-line-length = "128"
 
 [tool.pylint.messages_control]
-disable = "C0114"
+disable = "C0114,R0902,R0913"
 
 [tool.ruff]
 select = ["ALL"]
@@ -83,8 +83,9 @@ ignore = [
     "ANN401", # any-type
     "FBT",    # flake8-boolean-trap
     "S101",   # assert-used
-    "COM812", #confilcts with ruff formatter
-    "ISC001", #confilcts with ruff formatter
+    "COM812", # confilcts with ruff formatter
+    "ISC001", # confilcts with ruff formatter
+    "PLR0913",# Too many arguments in function definition
     ]
 line-length = 128
 target-version = "py38"
diff --git a/src/simweights/_corsika_weighter.py b/src/simweights/_corsika_weighter.py
index 5c5e192..e90c969 100644
--- a/src/simweights/_corsika_weighter.py
+++ b/src/simweights/_corsika_weighter.py
@@ -13,7 +13,7 @@
 from ._generation_surface import GenerationSurface, generation_surface
 from ._powerlaw import PowerLaw
 from ._spatial import NaturalRateCylinder
-from ._utils import constcol, get_column, get_table, has_table
+from ._utils import Column, constcol, get_column, get_table, has_table
 from ._weighter import Weighter
 
 
@@ -32,17 +32,20 @@ def sframe_corsika_surface(table: Any, oversampling: bool) -> GenerationSurface:
             get_column(table, "cylinder_radius")[i],
             np.cos(get_column(table, "max_zenith")[i]),
             np.cos(get_column(table, "min_zenith")[i]),
+            "cos_zen",
         )
         spectrum = PowerLaw(
             get_column(table, "power_law_index")[i],
             get_column(table, "min_energy")[i],
             get_column(table, "max_energy")[i],
+            "energy",
         )
         oversampling_val = get_column(table, "oversampling")[i] if oversampling else 1
+        pdgid = int(get_column(table, "primary_type")[i])
         surfaces.append(
             get_column(table, "n_events")[i]
             * oversampling_val
-            * generation_surface(int(get_column(table, "primary_type")[i]), spectrum, spatial),
+            * generation_surface(pdgid, Column("event_weight"), spectrum, spatial),
         )
     retval = sum(surfaces)
     assert isinstance(retval, GenerationSurface)
@@ -61,6 +64,7 @@ def weight_map_corsika_surface(table: Any) -> GenerationSurface:
             constcol(table, "CylinderRadius", mask),
             np.cos(constcol(table, "ThetaMax", mask)),
             np.cos(constcol(table, "ThetaMin", mask)),
+            "cos_zen",
         )
 
         primary_spectral_index = round(constcol(table, "PrimarySpectralIndex", mask), 6)
@@ -70,6 +74,7 @@ def weight_map_corsika_surface(table: Any) -> GenerationSurface:
             primary_spectral_index,
             constcol(table, "EnergyPrimaryMin", mask),
             constcol(table, "EnergyPrimaryMax", mask),
+            "energy",
         )
         nevents = constcol(table, "OverSampling", mask) * constcol(table, "NEvents", mask)
         surface += nevents * generation_surface(pdgid, spectrum, spatial)
diff --git a/src/simweights/_generation_surface.py b/src/simweights/_generation_surface.py
index 7ffa71d..6eb4e90 100644
--- a/src/simweights/_generation_surface.py
+++ b/src/simweights/_generation_surface.py
@@ -5,17 +5,20 @@
 from __future__ import annotations
 
 from copy import deepcopy
-from typing import TYPE_CHECKING, NamedTuple
+from typing import TYPE_CHECKING, NamedTuple, Sequence
 
 import numpy as np
 
 from ._pdgcode import PDGCode
+from ._powerlaw import PowerLaw
+from ._spatial import CircleInjector, CylinderBase, SpatialDist
 
-if TYPE_CHECKING:
+if TYPE_CHECKING:  # pragma: no cover
     from numpy.typing import ArrayLike, NDArray
 
-    from ._powerlaw import PowerLaw  # pragma: no cover
-    from ._spatial import SpatialDist  # pragma: no cover
+    from ._utils import Column, Const
+
+    Dist = SpatialDist | PowerLaw | Column | Const
 
 
 class SurfaceTuple(NamedTuple):
@@ -23,8 +26,7 @@ class SurfaceTuple(NamedTuple):
 
     pdgid: int | PDGCode
     nevents: float
-    energy_dist: PowerLaw
-    spatial_dist: SpatialDist
+    dists: Sequence
 
 
 class GenerationSurface:
@@ -46,7 +48,7 @@ def _insert(self: GenerationSurface, surface: SurfaceTuple) -> None:
             self.spectra[key] = []
 
         for i, spec in enumerate(self.spectra[key]):
-            if surface.energy_dist == spec.energy_dist and surface.spatial_dist == spec.spatial_dist:
+            if surface.dists == spec.dists:
                 self.spectra[key][i] = spec._replace(nevents=surface.nevents + spec.nevents)
                 break
         else:
@@ -77,30 +79,45 @@ def __mul__(self: GenerationSurface, factor: float) -> GenerationSurface:
     def __rmul__(self: GenerationSurface, factor: float) -> GenerationSurface:
         return self.__mul__(factor)
 
+    def get_keys(self: GenerationSurface) -> list[str]:
+        """Get a list of the available keys needed for weighting this surface."""
+        keys = []
+        for x in self.spectra.values():
+            for y in x:
+                for a in y.dists:
+                    keys += list(a.columns)
+        return list(set(keys))
+
     def get_epdf(
         self: GenerationSurface,
         pdgid: ArrayLike,
-        energy: ArrayLike,
-        cos_zen: ArrayLike,
+        **kwargs: ArrayLike,
     ) -> NDArray[np.float64]:
         """Get the extended pdf of an event.
 
         The pdf is the probability that an event with these parameters is generated. The pdf is multiplied
         by the number of events.
         """
-        energy = np.asarray(energy)
-        cos_zen = np.asarray(cos_zen)
-        count = np.zeros_like(energy, dtype=float)
-
+        cols = {}
+        shape = None
+        for key, value in kwargs.items():
+            cols[key] = np.asfarray(value)
+            if shape is None:
+                shape = cols[key].shape
+            else:
+                assert shape == cols[key].shape  # type: ignore[unreachable]
+        assert shape is not None
+        count = np.zeros(shape, dtype=np.float_)
+        # loop over particle type
         for ptype in np.unique(pdgid):
             mask = ptype == pdgid
-            if np.any(mask):
-                masked_energy = energy[mask]
-                masked_cos_zen = cos_zen[mask]
-                count[mask] += sum(
-                    p.nevents * p.energy_dist.pdf(masked_energy) * p.spatial_dist.pdf(masked_cos_zen)
-                    for p in self.spectra[ptype]
-                )
+            # loop over different datasets of the same particle type
+            for surface in self.spectra[ptype]:
+                result = surface.nevents
+                # loop over the different distributions in the generation surface
+                for dist in surface.dists:
+                    result *= dist.pdf(*(cols[k][mask] for k in dist.columns))
+                count[mask] += result
         return count
 
     def get_pdgids(self: GenerationSurface) -> list[int | PDGCode]:
@@ -116,8 +133,10 @@ def get_energy_range(self: GenerationSurface, pdgid: PDGCode | None) -> tuple[fl
         emax = -np.inf
         for pid in pdgids:
             for surf in self.spectra[pid]:
-                emin = min(emin, surf.energy_dist.a)
-                emax = max(emax, surf.energy_dist.b)
+                for dist in surf.dists:
+                    if isinstance(dist, PowerLaw):
+                        emin = min(emin, dist.a)
+                        emax = max(emax, dist.b)
         assert np.isfinite(emin)
         assert np.isfinite(emax)
         return emin, emax
@@ -131,8 +150,10 @@ def get_cos_zenith_range(self: GenerationSurface, pdgid: PDGCode | None) -> tupl
         czmax = -np.inf
         for pid in pdgids:
             for surf in self.spectra[pid]:
-                czmin = min(czmin, surf.spatial_dist.cos_zen_min)
-                czmax = max(czmax, surf.spatial_dist.cos_zen_max)
+                for dist in surf.dists:
+                    if isinstance(dist, (CircleInjector, CylinderBase)):
+                        czmin = min(czmin, dist.cos_zen_min)
+                        czmax = max(czmax, dist.cos_zen_max)
         assert np.isfinite(czmin)
         assert np.isfinite(czmax)
         return czmin, czmax
@@ -164,18 +185,13 @@ def __str__(self: GenerationSurface) -> str:
                 ptype = PDGCode(pdgid).name
             except ValueError:
                 ptype = str(pdgid)
-
-            collections = [f"N={subspec.nevents} {subspec.energy_dist} {subspec.spatial_dist}" for subspec in specs]
+            collections = (f"N={subspec.nevents} " + " ".join(repr(d) for d in subspec.dists) for subspec in specs)
             outstrs.append(f"     {ptype:>11} : " + "\n                   ".join(collections))
         return "< " + self.__class__.__name__ + "\n" + "\n".join(outstrs) + "\n>"
 
 
-def generation_surface(
-    pdgid: int | PDGCode,
-    energy_dist: PowerLaw,
-    spatial_dist: SpatialDist,
-) -> GenerationSurface:
+def generation_surface(pdgid: int | PDGCode, *dists: Dist) -> GenerationSurface:
     """Convenience function to generate a GenerationSurface for a single particle type."""
     return GenerationSurface(
-        SurfaceTuple(pdgid=pdgid, nevents=1.0, energy_dist=energy_dist, spatial_dist=spatial_dist),
+        SurfaceTuple(pdgid=pdgid, nevents=1.0, dists=dists),
     )
diff --git a/src/simweights/_genie_weighter.py b/src/simweights/_genie_weighter.py
index d3ae17b..75d668c 100644
--- a/src/simweights/_genie_weighter.py
+++ b/src/simweights/_genie_weighter.py
@@ -10,7 +10,7 @@
 from ._generation_surface import GenerationSurface, generation_surface
 from ._powerlaw import PowerLaw
 from ._spatial import CircleInjector
-from ._utils import constcol, get_column, get_table
+from ._utils import Column, Const, get_column, get_table
 from ._weighter import Weighter
 
 
@@ -29,10 +29,14 @@ def genie_surface(table: Iterable[Mapping[str, float]]) -> GenerationSurface:
             -get_column(table, "power_law_index")[i],
             get_column(table, "min_energy")[i],
             get_column(table, "max_energy")[i],
+            "energy",
         )
+        pdgid = int(get_column(table, "primary_type")[i])
+        nevents = get_column(table, "n_flux_events")[i]
+        global_probability_scale = get_column(table, "global_probability_scale")[i]
         surfaces.append(
-            get_column(table, "n_flux_events")[i]
-            * generation_surface(int(get_column(table, "primary_type")[i]), spectrum, spatial),
+            nevents
+            * generation_surface(pdgid, Const(1 / spatial.etendue / global_probability_scale), Column("wght"), spectrum),
         )
     retval = sum(surfaces)
     assert isinstance(retval, GenerationSurface)
@@ -47,15 +51,10 @@ def GenieWeighter(file_obj: Any) -> Weighter:  # noqa: N802
     """
     weight_table = get_table(file_obj, "I3GenieInfo")
     surface = genie_surface(weight_table)
-    global_probability_scale = constcol(weight_table, "global_probability_scale")
 
     weighter = Weighter([file_obj], surface)
-    weighter.add_weight_column("energy", weighter.get_column("I3GenieResult", "Ev"))
     weighter.add_weight_column("pdgid", weighter.get_column("I3GenieResult", "neu").astype(np.int32))
-    weighter.add_weight_column("cos_zen", np.full(len(weighter.get_column("I3GenieResult", "Ev")), 1))
-    weighter.add_weight_column(
-        "event_weight",
-        global_probability_scale * weighter.get_column("I3GenieResult", "wght"),
-    )
+    weighter.add_weight_column("energy", weighter.get_column("I3GenieResult", "Ev"))
+    weighter.add_weight_column("wght", weighter.get_column("I3GenieResult", "wght"))
 
     return weighter
diff --git a/src/simweights/_icetop_weighter.py b/src/simweights/_icetop_weighter.py
index e95fb5c..fffd611 100644
--- a/src/simweights/_icetop_weighter.py
+++ b/src/simweights/_icetop_weighter.py
@@ -20,20 +20,22 @@ def sframe_icetop_surface(table: Any) -> GenerationSurface:
 
     for i in range(len(get_column(table, "n_events"))):
         assert get_column(table, "power_law_index")[i] <= 0
-        spectrum = PowerLaw(
+        spectrum = PowerLaw(  # pylint: disable=duplicate-code
             get_column(table, "power_law_index")[i],
             get_column(table, "min_energy")[i],
             get_column(table, "max_energy")[i],
+            "energy",
         )
         spatial = NaturalRateCylinder(
             0,  # set cylinder height to 0 to get simple surface plane
             get_column(table, "sampling_radius")[i],
             np.cos(get_column(table, "max_zenith")[i]),
             np.cos(get_column(table, "min_zenith")[i]),
+            "cos_zen",
         )
-        surfaces.append(
-            get_column(table, "n_events")[i] * generation_surface(int(get_column(table, "primary_type")[i]), spectrum, spatial),
-        )
+        nevents = get_column(table, "n_events")[i]
+        pdgid = int(get_column(table, "primary_type")[i])
+        surfaces.append(nevents * generation_surface(pdgid, spectrum, spatial))
     retval = sum(surfaces)
     assert isinstance(retval, GenerationSurface)
     return retval
@@ -49,6 +51,4 @@ def IceTopWeighter(file_obj: Any) -> Weighter:  # noqa: N802
     weighter.add_weight_column("pdgid", pdgid)
     weighter.add_weight_column("energy", weighter.get_column("MCPrimary", "energy"))
     weighter.add_weight_column("cos_zen", np.cos(weighter.get_column("MCPrimary", "zenith")))
-    weighter.add_weight_column("event_weight", np.ones_like(pdgid))
-
     return weighter
diff --git a/src/simweights/_nugen_weighter.py b/src/simweights/_nugen_weighter.py
index 29f0f0f..b82b04d 100644
--- a/src/simweights/_nugen_weighter.py
+++ b/src/simweights/_nugen_weighter.py
@@ -10,7 +10,7 @@
 from ._generation_surface import GenerationSurface, generation_surface
 from ._powerlaw import PowerLaw
 from ._spatial import CircleInjector, SpatialDist, UniformSolidAngleCylinder
-from ._utils import constcol, get_column, get_table, has_column
+from ._utils import Column, constcol, get_column, get_table, has_column
 from ._weighter import Weighter
 
 
@@ -30,14 +30,14 @@ def nugen_spatial(table: Any) -> SpatialDist:
     injection_radius = constcol(table, "InjectionSurfaceR") if has_column(table, "InjectionSurfaceR") else -1
 
     if injection_radius > 0:
-        return CircleInjector(injection_radius, min_cos, max_cos)
+        return CircleInjector(injection_radius, min_cos, max_cos, "cos_zen")
 
     # Surface mode was added in V04-01-00 but the cylinder size was hard coded, `CylinderHeight` and
     # `CylinderRadius` were added after later V06-00-00. If they are not in the table then use the
     # hardcoded values
     cylinder_height = constcol(table, "CylinderHeight") if has_column(table, "CylinderHeight") else 1900
     cylinder_radius = constcol(table, "CylinderRadius") if has_column(table, "CylinderRadius") else 950
-    return UniformSolidAngleCylinder(cylinder_height, cylinder_radius, min_cos, max_cos)
+    return UniformSolidAngleCylinder(cylinder_height, cylinder_radius, min_cos, max_cos, "cos_zen")
 
 
 def nugen_spectrum(table: Any) -> PowerLaw:
@@ -48,7 +48,7 @@ def nugen_spectrum(table: Any) -> PowerLaw:
     # for negative slopes ie +2 means E**-2 injection spectrum
     power_law_index = -constcol(table, "PowerLawIndex")
     assert power_law_index <= 0
-    return PowerLaw(power_law_index, min_energy, max_energy)
+    return PowerLaw(power_law_index, min_energy, max_energy, "energy")
 
 
 def nugen_surface(table: Any) -> GenerationSurface:
@@ -66,9 +66,9 @@ def nugen_surface(table: Any) -> GenerationSurface:
         # newer version will explicitly put the ratio in `TypeWeight` but for older version we
         # assume it is 0.5
         type_weight = constcol(table, "TypeWeight", mask) if has_column(table, "TypeWeight") else 0.5
-        primary_type = constcol(table, "PrimaryNeutrinoType", mask)
+        primary_type = int(constcol(table, "PrimaryNeutrinoType", mask))
         n_events = type_weight * constcol(table, "NEvents", mask)
-        surfaces.append(n_events * generation_surface(int(primary_type), spectrum, spatial))
+        surfaces.append(n_events * generation_surface(primary_type, Column("event_weight"), spectrum, spatial))
     ret = sum(surfaces)
     assert isinstance(ret, GenerationSurface)
     return ret
diff --git a/src/simweights/_powerlaw.py b/src/simweights/_powerlaw.py
index 86de033..16e78f4 100644
--- a/src/simweights/_powerlaw.py
+++ b/src/simweights/_powerlaw.py
@@ -38,7 +38,7 @@ class PowerLaw:
 
     # pylint: disable=invalid-name
 
-    def __init__(self: PowerLaw, g: float, a: float, b: float) -> None:
+    def __init__(self: PowerLaw, g: float, a: float, b: float, colname: str | None = None) -> None:
         assert b > a
         self.g = float(g)
         self.a = float(a)
@@ -48,8 +48,8 @@ def __init__(self: PowerLaw, g: float, a: float, b: float) -> None:
             self.integral = np.log(self.b / self.a)
         else:
             self.integral = (self.b**self.G - self.a**self.G) / self.G
-
         self.span = b - a
+        self.columns = (colname,)
 
     def _pdf(self: PowerLaw, x: NDArray[np.float64]) -> NDArray[np.float64]:
         return np.asfarray(x**self.g / self.integral)
diff --git a/src/simweights/_spatial.py b/src/simweights/_spatial.py
index 39d6857..3ecc6a1 100644
--- a/src/simweights/_spatial.py
+++ b/src/simweights/_spatial.py
@@ -20,6 +20,7 @@ def __init__(
         radius: float,
         cos_zen_min: float,
         cos_zen_max: float,
+        colname: str | None = None,
     ) -> None:
         if cos_zen_min < -1 or cos_zen_max > 1:
             raise ValueError(
@@ -34,6 +35,7 @@ def __init__(
         self._side = 2e4 * self.radius * self.length
         self._cap = 1e4 * np.pi * self.radius**2
         self.etendue = float(self._diff_etendue(self.cos_zen_max) - self._diff_etendue(self.cos_zen_min))
+        self.columns = (colname,)
 
     def projected_area(self: CylinderBase, cos_zen: ArrayLike) -> NDArray[np.float64]:
         """Cross sectional area of a cylinder in cm^2.
@@ -109,8 +111,15 @@ class NaturalRateCylinder(CylinderBase):
         I \propto \pi\cdot r^2\cdot\sin\theta\cdot(\cos\theta+2/\pi\cdot l/r\cdot\sin\theta)
     """
 
-    def __init__(self: NaturalRateCylinder, length: float, radius: float, cos_zen_min: float, cos_zen_max: float) -> None:
-        super().__init__(length, radius, cos_zen_min, cos_zen_max)
+    def __init__(
+        self: NaturalRateCylinder,
+        length: float,
+        radius: float,
+        cos_zen_min: float,
+        cos_zen_max: float,
+        colname: str | None = None,
+    ) -> None:
+        super().__init__(length, radius, cos_zen_min, cos_zen_max, colname)
         self._normalization = 1 / self.etendue
 
     def pdf(self: NaturalRateCylinder, cos_zen: ArrayLike) -> NDArray[np.float64]:
@@ -130,13 +139,16 @@ class CircleInjector:
     The etendue is just the area of the circle times the solid angle.
     """
 
-    def __init__(self: CircleInjector, radius: float, cos_zen_min: float, cos_zen_max: float) -> None:
+    def __init__(
+        self: CircleInjector, radius: float, cos_zen_min: float, cos_zen_max: float, colname: str | None = None
+    ) -> None:
         self.radius = radius
         self.cos_zen_min = cos_zen_min
         self.cos_zen_max = cos_zen_max
         self._cap = 1e4 * np.pi * self.radius**2
         self.etendue = 2 * np.pi * (self.cos_zen_max - self.cos_zen_min) * self._cap
         self._normalization = 1 / self.etendue
+        self.columns = (colname,)
 
     def projected_area(self: CircleInjector, cos_zen: float) -> float:  # noqa: ARG002
         """Returns the cross sectional area of the injection area in cm^2."""
diff --git a/src/simweights/_utils.py b/src/simweights/_utils.py
index 6aebccb..320fac9 100644
--- a/src/simweights/_utils.py
+++ b/src/simweights/_utils.py
@@ -12,7 +12,7 @@
 
 from ._pdgcode import PDGCode
 
-if TYPE_CHECKING:
+if TYPE_CHECKING:  # pragma: no cover
     from numpy.typing import ArrayLike, NDArray
 
 
@@ -21,6 +21,44 @@
 SeedType = Union[GeneratorType, IntNumber, None]
 
 
+class Column:
+    """Simple PDF class for a pdf that just uses a column.
+
+    Usually used for stuff like probability of interaction.
+    """
+
+    def __init__(self: Column, colname: str | None = None) -> None:
+        self.columns = (colname,)
+
+    def pdf(self: Column, value: ArrayLike) -> NDArray[np.float64]:
+        r"""Probability density function."""
+        return 1 / np.asfarray(value)
+
+    def __eq__(self: Column, other: object) -> bool:
+        return isinstance(other, Column) and self.columns == other.columns
+
+    def __str__(self: Column) -> str:
+        return f"Column{self.columns!r}"
+
+
+class Const:
+    """Simple PDF class for a supplied constant."""
+
+    def __init__(self: Const, v: float) -> None:
+        self.columns = ()
+        self.v = v
+
+    def pdf(self: Const) -> NDArray[np.float64]:
+        r"""Probability density function."""
+        return np.asfarray(self.v)
+
+    def __eq__(self: Const, other: object) -> bool:
+        return isinstance(other, Const) and self.v == other.v
+
+    def __str__(self: Const) -> str:
+        return f"Const({self.v})"
+
+
 def has_table(file_obj: Any, name: str) -> bool:
     """Helper function for determining if a file has a table, works with h5py, pytables, and pandas."""
     if hasattr(file_obj, "root"):
diff --git a/src/simweights/_weighter.py b/src/simweights/_weighter.py
index 7bfd256..36c76db 100644
--- a/src/simweights/_weighter.py
+++ b/src/simweights/_weighter.py
@@ -79,14 +79,14 @@ def get_weights(self: Weighter, flux: Any) -> NDArray[np.float64]:
             areas. For neutrinos, If the value is 1 then the return value will be the
             well known quantity OneWeight.
         """
-        event_col = {k: self.get_weight_column(k) for k in ["energy", "pdgid", "cos_zen"]}
+        # get a dictionary of the columns we need for weighting
+        event_col = {k: self.get_weight_column(k) for k in ["pdgid", *self.surface.get_keys()]}
 
         # do nothing if everything is empty
         if event_col["pdgid"].shape == (0,):
             return np.array([])
 
         epdf = self.surface.get_epdf(**event_col)
-        event_weight = self.get_weight_column("event_weight")
 
         # calculate the flux based on which type of flux it is
         if hasattr(flux, "getFlux"):
@@ -132,7 +132,7 @@ def get_weights(self: Weighter, flux: Any) -> NDArray[np.float64]:
                 stacklevel=2,
             )
         weights = np.zeros_like(epdf)
-        weights[mask] = (event_weight * flux_val)[mask] / epdf[mask]
+        weights[mask] = flux_val[mask] / epdf[mask]
         return weights
 
     def effective_area(
diff --git a/tests/test_generation_surface.py b/tests/test_generation_surface.py
index 942247b..e531adc 100755
--- a/tests/test_generation_surface.py
+++ b/tests/test_generation_surface.py
@@ -8,6 +8,7 @@
 from copy import deepcopy
 
 import numpy as np
+from numpy.testing import assert_allclose
 from simweights import GenerationSurface, NaturalRateCylinder, PDGCode, PowerLaw, generation_surface
 from simweights._generation_surface import SurfaceTuple  # noqa: F401
 
@@ -15,10 +16,10 @@
 class Testgeneration_surface(unittest.TestCase):
     @classmethod
     def setUpClass(cls):
-        cls.p1 = PowerLaw(-1, 10, 100)
-        cls.p2 = PowerLaw(-2, 50, 500)
-        cls.c1 = NaturalRateCylinder(3, 8, -1, 1)
-        cls.c2 = NaturalRateCylinder(4, 8, -1, 1)
+        cls.p1 = PowerLaw(-1, 10, 100, "energy")
+        cls.p2 = PowerLaw(-2, 50, 500, "energy")
+        cls.c1 = NaturalRateCylinder(3, 8, -1, 1, "cos_zen")
+        cls.c2 = NaturalRateCylinder(4, 8, -1, 1, "cos_zen")
 
         cls.N1 = 10000
         cls.N2 = 20000
@@ -108,9 +109,9 @@ def test_addition(self):
         self.assertNotEqual(s, self.s1)
         self.assertEqual(n0, self.s0.spectra[2212][0].nevents)
         self.assertEqual(n1, self.s1.spectra[2212][0].nevents)
-        self.assertAlmostEqual(
-            s.get_epdf(2212, 50, 0),
-            self.s0.get_epdf(2212, 50, 0) + self.s1.get_epdf(2212, 50, 0),
+        assert_allclose(
+            s.get_epdf(2212, energy=50, cos_zen=0),
+            self.s0.get_epdf(2212, energy=50, cos_zen=0) + self.s1.get_epdf(2212, energy=50, cos_zen=0),
         )
 
         ss = self.s0 + self.s2
@@ -119,9 +120,10 @@ def test_addition(self):
         self.assertEqual(len(ss.spectra[2212]), 2)
         self.assertEqual(ss.spectra[2212][0], self.s0.spectra[2212][0])
         self.assertEqual(ss.spectra[2212][1], self.s2.spectra[2212][0])
+
         self.assertAlmostEqual(
-            ss.get_epdf(2212, 50, 0),
-            self.s0.get_epdf(2212, 50, 0) + self.s2.get_epdf(2212, 50, 0),
+            ss.get_epdf(2212, energy=50, cos_zen=0),
+            self.s0.get_epdf(2212, energy=50, cos_zen=0) + self.s2.get_epdf(2212, energy=50, cos_zen=0),
         )
 
         s3 = self.s0 + self.s3
@@ -131,8 +133,8 @@ def test_addition(self):
         self.assertEqual(s3.spectra[2212][0], self.s0.spectra[2212][0])
         self.assertEqual(s3.spectra[2212][1], self.s3.spectra[2212][0])
         self.assertAlmostEqual(
-            s3.get_epdf(2212, 50, 0),
-            self.s0.get_epdf(2212, 50, 0) + self.s3.get_epdf(2212, 50, 0),
+            s3.get_epdf(2212, energy=50, cos_zen=0),
+            self.s0.get_epdf(2212, energy=50, cos_zen=0) + self.s3.get_epdf(2212, energy=50, cos_zen=0),
         )
 
         s4 = self.s0 + self.s4
@@ -142,8 +144,8 @@ def test_addition(self):
         self.assertEqual(len(s4.spectra[2213]), 1)
         self.assertEqual(s4.spectra[2212][0], self.s0.spectra[2212][0])
         self.assertEqual(s4.spectra[2213][0], self.s4.spectra[2213][0])
-        self.assertAlmostEqual(s4.get_epdf(2212, 50, 0), self.s0.get_epdf(2212, 50, 0))
-        self.assertAlmostEqual(s4.get_epdf(2213, 50, 0), self.s4.get_epdf(2213, 50, 0))
+        self.assertAlmostEqual(s4.get_epdf(2212, energy=50, cos_zen=0), self.s0.get_epdf(2212, energy=50, cos_zen=0))
+        self.assertAlmostEqual(s4.get_epdf(2213, energy=50, cos_zen=0), self.s4.get_epdf(2213, energy=50, cos_zen=0))
 
         with self.assertRaises(TypeError):
             self.s0 + 47
@@ -168,19 +170,19 @@ def test_multiplication(self):
         sa *= 4.4
         self.assertEqual(sa.spectra[2212][0].nevents, 44000)
         self.assertEqual(self.s0.spectra[2212][0].nevents, 10000)
-        self.assertAlmostEqual(sa.get_epdf(2212, 50, 0), 4.4 * self.s0.get_epdf(2212, 50, 0))
+        assert_allclose(sa.get_epdf(2212, energy=50, cos_zen=0), 4.4 * self.s0.get_epdf(2212, energy=50, cos_zen=0))
 
         sb = self.s0 * 5.5
         self.assertNotEqual(id(sb), id(self.s0))
         self.assertEqual(sb.spectra[2212][0].nevents, 55000)
         self.assertEqual(self.s0.spectra[2212][0].nevents, 10000)
-        self.assertAlmostEqual(sb.get_epdf(2212, 50, 0), 5.5 * self.s0.get_epdf(2212, 50, 0))
+        assert_allclose(sb.get_epdf(2212, energy=50, cos_zen=0), 5.5 * self.s0.get_epdf(2212, energy=50, cos_zen=0))
 
         sc = 6.6 * self.s0
         self.assertNotEqual(id(sc), id(self.s0))
         self.assertEqual(sc.spectra[2212][0].nevents, 66000)
         self.assertEqual(self.s0.spectra[2212][0].nevents, 10000)
-        self.assertAlmostEqual(sc.get_epdf(2212, 50, 0), 6.6 * self.s0.get_epdf(2212, 50, 0))
+        assert_allclose(sc.get_epdf(2212, energy=50, cos_zen=0), 6.6 * self.s0.get_epdf(2212, energy=50, cos_zen=0))
 
     def test_repr(self):
         PPlus = PDGCode.PPlus  # noqa: F841
@@ -194,20 +196,20 @@ def test_powerlaw(self):
         N = int(self.s0.spectra[2212][0].nevents)
         E = np.geomspace(self.p1.a, self.p1.b - 1 / N, N)
         cz = np.linspace(self.c1.cos_zen_min, self.c1.cos_zen_max, N)
-        w = 1 / self.s0.get_epdf(2212, E, cz)
+        w = 1 / self.s0.get_epdf(2212, energy=E, cos_zen=cz)
 
         area = (self.p1.b - self.p1.a) * (2e4 * self.c1.radius * np.pi**2 * (self.c1.radius + self.c1.length))
 
         self.assertAlmostEqual(
             area,
-            self.s0.spectra[2212][0].energy_dist.span * self.s0.spectra[2212][0].spatial_dist.etendue,
+            self.s0.spectra[2212][0].dists[0].span * self.s0.spectra[2212][0].dists[1].etendue,
         )
         self.assertAlmostEqual(w.sum() / area, 1, 4)
 
-        self.assertEqual(self.s0.spectra[2212][0].spatial_dist, self.c1)
-        self.assertIsNot(self.s0.spectra[2212][0].spatial_dist, self.c1)
-        self.assertEqual(self.s0.spectra[2212][0].energy_dist, self.p1)
-        self.assertIsNot(self.s0.spectra[2212][0].energy_dist, self.p1)
+        self.assertEqual(self.s0.spectra[2212][0].dists[1], self.c1)
+        self.assertIsNot(self.s0.spectra[2212][0].dists[1], self.c1)
+        self.assertEqual(self.s0.spectra[2212][0].dists[0], self.p1)
+        self.assertIsNot(self.s0.spectra[2212][0].dists[0], self.p1)
 
     def test_two_surfaces(self):
         N = int(self.s0.spectra[2212][0].nevents)
@@ -219,36 +221,36 @@ def test_two_surfaces(self):
         surf = self.s0 + self.s2
         E = np.r_[E1, E2]
         czc = np.r_[cz, cz]
-        wc = 1 / surf.get_epdf(2212, E, czc)
+        wc = 1 / surf.get_epdf(2212, energy=E, cos_zen=czc)
 
         self.assertAlmostEqual(wc.sum() / (self.p2.b - self.p1.a) / self.c1.etendue, 1, 3)
 
-        self.assertEqual(self.s0.spectra[2212][0].spatial_dist, self.c1)
-        self.assertIsNot(self.s0.spectra[2212][0].spatial_dist, self.c1)
-        self.assertEqual(self.s0.spectra[2212][0].energy_dist, self.p1)
-        self.assertIsNot(self.s0.spectra[2212][0].energy_dist, self.p1)
+        self.assertEqual(self.s0.spectra[2212][0].dists[1], self.c1)
+        self.assertIsNot(self.s0.spectra[2212][0].dists[1], self.c1)
+        self.assertEqual(self.s0.spectra[2212][0].dists[0], self.p1)
+        self.assertIsNot(self.s0.spectra[2212][0].dists[0], self.p1)
 
-        self.assertEqual(self.s2.spectra[2212][0].spatial_dist, self.c1)
-        self.assertIsNot(self.s2.spectra[2212][0].spatial_dist, self.c1)
-        self.assertEqual(self.s2.spectra[2212][0].energy_dist, self.p2)
-        self.assertIsNot(self.s2.spectra[2212][0].energy_dist, self.p2)
+        self.assertEqual(self.s2.spectra[2212][0].dists[1], self.c1)
+        self.assertIsNot(self.s2.spectra[2212][0].dists[1], self.c1)
+        self.assertEqual(self.s2.spectra[2212][0].dists[0], self.p2)
+        self.assertIsNot(self.s2.spectra[2212][0].dists[0], self.p2)
 
         self.assertEqual(len(surf.spectra), 1)
         np.testing.assert_array_equal(list(surf.spectra.keys()), [2212])
 
         self.assertEqual(surf.spectra[2212][0], self.s0.spectra[2212][0])
         self.assertIsNot(surf.spectra[2212][0], self.s0.spectra[2212][0])
-        self.assertEqual(surf.spectra[2212][0].spatial_dist, self.s0.spectra[2212][0].spatial_dist)
-        self.assertIsNot(surf.spectra[2212][0].spatial_dist, self.s0.spectra[2212][0].spatial_dist)
-        self.assertEqual(surf.spectra[2212][0].energy_dist, self.s0.spectra[2212][0].energy_dist)
-        self.assertIsNot(surf.spectra[2212][0].energy_dist, self.s0.spectra[2212][0].energy_dist)
+        self.assertEqual(surf.spectra[2212][0].dists[1], self.s0.spectra[2212][0].dists[1])
+        self.assertIsNot(surf.spectra[2212][0].dists[1], self.s0.spectra[2212][0].dists[1])
+        self.assertEqual(surf.spectra[2212][0].dists[0], self.s0.spectra[2212][0].dists[0])
+        self.assertIsNot(surf.spectra[2212][0].dists[0], self.s0.spectra[2212][0].dists[0])
 
         self.assertEqual(surf.spectra[2212][1], self.s2.spectra[2212][0])
         self.assertIsNot(surf.spectra[2212][1], self.s2.spectra[2212][0])
-        self.assertEqual(surf.spectra[2212][1].spatial_dist, self.s2.spectra[2212][0].spatial_dist)
-        self.assertIsNot(surf.spectra[2212][1].spatial_dist, self.s2.spectra[2212][0].spatial_dist)
-        self.assertEqual(surf.spectra[2212][1].energy_dist, self.s2.spectra[2212][0].energy_dist)
-        self.assertIsNot(surf.spectra[2212][1].energy_dist, self.s2.spectra[2212][0].energy_dist)
+        self.assertEqual(surf.spectra[2212][1].dists[1], self.s2.spectra[2212][0].dists[1])
+        self.assertIsNot(surf.spectra[2212][1].dists[1], self.s2.spectra[2212][0].dists[1])
+        self.assertEqual(surf.spectra[2212][1].dists[0], self.s2.spectra[2212][0].dists[0])
+        self.assertIsNot(surf.spectra[2212][1].dists[0], self.s2.spectra[2212][0].dists[0])
 
     def test_addition_gsc(self):
         sa = self.gsc1 + 0
diff --git a/tests/test_genie_datasets.py b/tests/test_genie_datasets.py
index 2ca69a5..1c65b74 100755
--- a/tests/test_genie_datasets.py
+++ b/tests/test_genie_datasets.py
@@ -32,12 +32,13 @@ def cmp_dataset(self, fname):
 
         solid_angle = 2 * np.pi * (np.cos(wd["MinZenith"]) - np.cos(wd["MaxZenith"]))
         injection_area = np.pi * (wd["InjectionSurfaceR"] * 1e2) ** 2
-        total_prob = wd["TotalInteractionProbabilityWeight"]
+        global_probability_scale = wd["GlobalProbabilityScale"]
+        genie_weight = wd["GENIEWeight"]
 
         pli = -wd["PowerLawIndex"][0]
         energy_integral = ((10 ** wd["MaxEnergyLog"][0]) ** (pli + 1) - (10 ** wd["MinEnergyLog"][0]) ** (pli + 1)) / (pli + 1)
         energy_factor = 1 / (wd["PrimaryNeutrinoEnergy"] ** pli / energy_integral)
-        one_weight = total_prob * energy_factor * solid_angle * injection_area
+        one_weight = global_probability_scale * genie_weight * energy_factor * solid_angle * injection_area
         np.testing.assert_allclose(one_weight, wd["OneWeight"])
         final_weight = wd["OneWeight"] / (get_column(get_table(reffile, "I3GenieInfo"), "n_flux_events")[0])
 
@@ -52,20 +53,16 @@ def cmp_dataset(self, fname):
             with self.subTest(lib=str(fobj)):
                 w = GenieWeighter(fobj)
 
-                np.testing.assert_allclose(w.get_weight_column("event_weight"), total_prob)
+                pdf0 = w.surface.spectra[14][0].dists[0]
+                np.testing.assert_allclose(1 / pdf0.v, global_probability_scale * solid_angle * injection_area, 1e-5)
 
-                cylinder = w.surface.spectra[14][0].spatial_dist
-                proj_area = cylinder.projected_area(0)
-                np.testing.assert_allclose(proj_area, injection_area)
+                np.testing.assert_allclose(w.get_weight_column("wght"), genie_weight)
 
-                sw_etendue = 1 / cylinder.pdf(0)
-                np.testing.assert_allclose(sw_etendue, solid_angle * injection_area, 1e-5)
-
-                power_law = w.surface.spectra[14][0].energy_dist
+                power_law = w.surface.spectra[14][0].dists[2]
                 energy_term = 1 / power_law.pdf(w.get_weight_column("energy"))
                 np.testing.assert_allclose(energy_term, energy_factor)
 
-                one_weight = w.get_weight_column("event_weight") * energy_term / cylinder.pdf(0)
+                one_weight = w.get_weight_column("wght") * energy_term * solid_angle * injection_area * global_probability_scale
                 np.testing.assert_allclose(one_weight, wd["OneWeight"], 1e-5)
 
                 np.testing.assert_allclose(w.get_weights(1), final_weight, 1e-5)
diff --git a/tests/test_icetop_datasets.py b/tests/test_icetop_datasets.py
index 4a14f7c..0b5b782 100755
--- a/tests/test_icetop_datasets.py
+++ b/tests/test_icetop_datasets.py
@@ -47,7 +47,7 @@ def cmp_dataset(self, fname):
         for fobj in fobjs:
             with self.subTest(lib=str(fobj)):
                 w = IceTopWeighter(fobj)
-                spatial = w.surface.spectra[2212][0].spatial_dist
+                spatial = w.surface.spectra[2212][0].dists[1]
                 proj_area = spatial.projected_area(1)
                 np.testing.assert_allclose(proj_area, injection_area)
                 sw_etendue = 1 / spatial.pdf(1)
diff --git a/tests/test_nugen_datasets.py b/tests/test_nugen_datasets.py
index 957e267..d5622fc 100755
--- a/tests/test_nugen_datasets.py
+++ b/tests/test_nugen_datasets.py
@@ -62,14 +62,14 @@ def cmp_dataset(self, fname):
                 event_weight = w.get_weight_column("event_weight")
                 np.testing.assert_allclose(event_weight, total_weight)
 
-                cylinder = w.surface.spectra[pdgid][0].spatial_dist
+                cylinder = w.surface.spectra[pdgid][0].dists[2]
                 proj_area = cylinder.projected_area(w.get_weight_column("cos_zen"))
                 np.testing.assert_allclose(proj_area, injection_area)
 
                 sw_etendue = 1 / cylinder.pdf(w.get_weight_column("cos_zen"))
                 np.testing.assert_allclose(sw_etendue, solid_angle * injection_area, 1e-5)
 
-                power_law = w.surface.spectra[pdgid][0].energy_dist
+                power_law = w.surface.spectra[pdgid][0].dists[1]
                 energy_factor = 1 / power_law.pdf(w.get_weight_column("energy"))
                 one_weight = w.get_weight_column("event_weight") * energy_factor * solid_angle * injection_area
                 np.testing.assert_allclose(one_weight, wd["OneWeight"])
diff --git a/tests/test_util.py b/tests/test_util.py
index d8497fa..1240e83 100755
--- a/tests/test_util.py
+++ b/tests/test_util.py
@@ -4,12 +4,25 @@
 #
 # SPDX-License-Identifier: BSD-2-Clause
 
+import tempfile
 import unittest
 from types import SimpleNamespace
 
 import numpy as np
+import uproot
+from numpy.testing import assert_array_equal
 from simweights import Hoerandel
-from simweights._utils import constcol, corsika_to_pdg, get_column, get_table, has_column, has_table
+from simweights._utils import (
+    Column,
+    Const,
+    check_random_state,
+    constcol,
+    corsika_to_pdg,
+    get_column,
+    get_table,
+    has_column,
+    has_table,
+)
 
 
 class TestUtil(unittest.TestCase):
@@ -34,16 +47,16 @@ def test_table_and_column(self):
         self.assertEqual(has_column(t1, "a"), True)
         self.assertEqual(has_column(t1, "b"), True)
         self.assertEqual(has_column(t1, "c"), False)
-        np.testing.assert_array_equal(get_column(t1, "a"), 10 * [3])
-        np.testing.assert_array_equal(get_column(t1, "b"), 5 * [3] + 5 * [4])
+        assert_array_equal(get_column(t1, "a"), 10 * [3])
+        assert_array_equal(get_column(t1, "b"), 5 * [3] + 5 * [4])
         with self.assertRaises(AttributeError):
             get_column(t1, "c")
 
         self.assertEqual(has_column(t2, "a"), True)
         self.assertEqual(has_column(t2, "b"), True)
         self.assertEqual(has_column(t2, "c"), False)
-        np.testing.assert_array_equal(get_column(t2, "a"), 10 * [7])
-        np.testing.assert_array_equal(get_column(t2, "b"), range(10))
+        assert_array_equal(get_column(t2, "a"), 10 * [7])
+        assert_array_equal(get_column(t2, "b"), range(10))
         with self.assertRaises(KeyError):
             get_column(t2, "c")
 
@@ -65,11 +78,57 @@ def test_table_and_column(self):
         with self.assertRaises(KeyError):
             constcol(t2, "c")
 
+    def test_dists(self):
+        p1 = Const(33)
+        p2 = Const(44)
+        self.assertEqual(p1.pdf(), 33)
+        self.assertEqual(p2.pdf(), 44)
+        self.assertEqual(p1, p1)
+        self.assertEqual(p2, p2)
+        self.assertNotEqual(p1, p2)
+        self.assertEqual(str(p1), "Const(33)")
+        self.assertEqual(str(p1), "Const(33)")
+
+        p1 = Column("energy")
+        p2 = Column("cos_zen")
+        assert_array_equal(p1.pdf(np.arange(1000)), 1 / np.arange(1000))
+        assert_array_equal(p2.pdf(1 / np.arange(10)), np.arange(10, dtype=float))
+        self.assertEqual(p1, p1)
+        self.assertEqual(p2, p2)
+        self.assertNotEqual(p1, p2)
+        self.assertEqual(str(p1), "Column('energy',)")
+        self.assertEqual(str(p2), "Column('cos_zen',)")
+
+    def test_uproot(self):
+        with tempfile.TemporaryFile() as f:
+            file = uproot.recreate(f)
+            a1 = np.arange(1000)
+            a2 = 2 * np.arange(1000)
+            file["tree1"] = {"branch1": a1, "branch2": a2}
+
+            file = uproot.open(f)
+            t = get_table(file, "tree1")
+            assert_array_equal(get_column(t, "branch1"), a1)
+            assert_array_equal(get_column(t, "branch2"), a2)
+
+    def test_check_random_state(self):
+        self.assertIsInstance(check_random_state(None), np.random.Generator)
+        self.assertIsInstance(check_random_state(3), np.random.Generator)
+        self.assertIsInstance(check_random_state(np.int16(3)), np.random.Generator)
+        r = np.random.RandomState()
+        self.assertEqual(check_random_state(r), r)
+        g = np.random.default_rng()
+        self.assertEqual(check_random_state(g), g)
+        with self.assertRaises(ValueError):
+            check_random_state(object())
+        with self.assertRaises(ValueError):
+            check_random_state(33.3)
+
     def test_corsika_to_pdg(self):
         c = [14, 402, 703, 904, 1105, 1206, 1407, 1608, 1909, 2010, 2311, 2412, 2713, 2814]
         c += [3115, 3216, 3517, 4018, 3919, 4020, 4521, 4822, 5123, 5224, 5525, 5626]
         pdgid = [int(i) for i in Hoerandel.pdgids]
-        np.testing.assert_array_equal(corsika_to_pdg(c), pdgid)
+        assert_array_equal(corsika_to_pdg(c), pdgid)
 
 
 if __name__ == "__main__":
diff --git a/tests/test_weighter.py b/tests/test_weighter.py
index 0ebc2b9..27f9651 100755
--- a/tests/test_weighter.py
+++ b/tests/test_weighter.py
@@ -39,8 +39,8 @@ def setUpClass(cls):
                 "zenith": np.full(cls.N1, np.pi / 4),
             },
         }
-        cls.c1 = NaturalRateCylinder(100, 200, 0, 1)
-        cls.p1 = PowerLaw(0, 5e5, 5e6)
+        cls.c1 = NaturalRateCylinder(100, 200, 0, 1, "cos_zen")
+        cls.p1 = PowerLaw(0, 5e5, 5e6, "energy")
         cls.s1 = cls.N1 * generation_surface(2212, cls.p1, cls.c1)
         cls.m1 = {
             "pdgid": ("I3Weight", "type"),
@@ -157,7 +157,7 @@ def test_weights(self):
 
     def test_empty(self):
         fake_file = {"I3Weight": {"energy": [], "type": [], "zenith": []}}
-        weighter = Weighter([fake_file], 0)
+        weighter = Weighter([fake_file], self.s1)
         weighter.add_weight_column("energy", np.array([]))
         weighter.add_weight_column("pdgid", np.array([]))
         weighter.add_weight_column("cos_zen", np.array([]))
@@ -289,7 +289,6 @@ def test_nuflux(self):
         weighter1.add_weight_column("pdgid", data1["I3Weight"]["type"])
         weighter1.add_weight_column("energy", data1["I3Weight"]["energy"])
         weighter1.add_weight_column("cos_zen", np.cos(data1["I3Weight"]["zenith"]))
-        weighter1.add_weight_column("event_weight", np.full(N1, 1))
 
         honda = nuflux.makeFlux("honda2006")
         w = weighter1.get_weights(honda)