From 9838469212b2b40418ec04bf5e226880e857a19a Mon Sep 17 00:00:00 2001
From: fjebaker <fergusbkr@gmail.com>
Date: Sat, 25 May 2024 23:28:41 +0100
Subject: [PATCH 1/6] feat(grouping): regroup with minimum counts

---
 src/datasets/spectraldata.jl | 11 ++++++++---
 src/datasets/spectrum.jl     | 25 +++++++++++++++++++++++++
 2 files changed, 33 insertions(+), 3 deletions(-)

diff --git a/src/datasets/spectraldata.jl b/src/datasets/spectraldata.jl
index 53eaf8cd..38426bcd 100644
--- a/src/datasets/spectraldata.jl
+++ b/src/datasets/spectraldata.jl
@@ -286,7 +286,12 @@ function regroup!(dataset::SpectralData, grouping; safety_copy = false)
     dataset
 end
 
-regroup!(dataset::SpectralData) = regroup!(dataset, dataset.spectrum.grouping)
+function regroup!(dataset::SpectralData; min_counts = nothing)
+    if !isnothing(min_counts)
+        group_min_counts!(dataset.spectrum, min_counts)
+    end
+    regroup!(dataset, dataset.spectrum.grouping)
+end
 
 function normalize!(dataset::SpectralData)
     ΔE = bin_widths(dataset)
@@ -445,8 +450,8 @@ macro _forward_SpectralData_api(args)
             layout::SpectralFitting.AbstractDataLayout,
             t::$(T),
         ) = SpectralFitting.objective_transformer(layout, getfield(t, $(field)))
-        SpectralFitting.regroup!(t::$(T), args...) =
-            SpectralFitting.regroup!(getfield(t, $(field)), args...)
+        SpectralFitting.regroup!(t::$(T), args...; kwargs...) =
+            SpectralFitting.regroup!(getfield(t, $(field)), args...; kwargs...)
         SpectralFitting.restrict_domain!(t::$(T), args...) =
             SpectralFitting.restrict_domain!(getfield(t, $(field)), args...)
         SpectralFitting.mask_energies!(t::$(T), args...) =
diff --git a/src/datasets/spectrum.jl b/src/datasets/spectrum.jl
index f6ee3505..9100e424 100644
--- a/src/datasets/spectrum.jl
+++ b/src/datasets/spectrum.jl
@@ -85,6 +85,31 @@ function regroup!(spectrum::Spectrum{T}, grouping) where {T}
     spectrum
 end
 
+function group_min_counts!(spectrum::Spectrum, min_counts::Int)
+    NEW_GRP = 1
+    CONTINUE_GRP = 0
+
+    function _counts(x)
+        if spectrum.units == u"counts"
+            convert(Int, x)
+        elseif spectrum.units == u"counts / s"
+            convert(Int, x * spectrum.exposure_time)
+        end
+    end
+
+    sum::Int = 0
+    for (i, f) in enumerate(spectrum.data)
+        c = _counts(f)
+        sum += c
+        if sum >= min_counts
+            spectrum.grouping[i] = NEW_GRP
+            sum = 0
+        else
+            spectrum.grouping[i] = CONTINUE_GRP
+        end
+    end
+end
+
 function Base.resize!(spectrum::Spectrum, n::Int)
     resize!(spectrum.channels, n)
     resize!(spectrum.data, n)

From 4de4a1f7b68da36b8f3c7ea0e6764360dd18ec36 Mon Sep 17 00:00:00 2001
From: fjebaker <fergusbkr@gmail.com>
Date: Sat, 25 May 2024 23:29:42 +0100
Subject: [PATCH 2/6] fix!: refactor AbstractDataLayout as AbstractLayout

---
 docs/src/datasets/datasets.md       |  2 +-
 docs/src/examples/optimizers.md     | 12 ++++++++++++
 docs/src/examples/sherpa-example.md |  2 +-
 docs/src/models/using-models.md     |  2 +-
 src/datasets/binning.jl             |  6 +++---
 src/datasets/datasets.jl            | 30 ++++++++++++++---------------
 src/datasets/injectivedata.jl       |  9 +++------
 src/datasets/spectraldata.jl        | 24 ++++++++++-------------
 src/fitting/cache.jl                |  8 ++++----
 src/support.jl                      | 26 ++++++++++++-------------
 10 files changed, 63 insertions(+), 58 deletions(-)

diff --git a/docs/src/datasets/datasets.md b/docs/src/datasets/datasets.md
index c6a99163..82a94b90 100644
--- a/docs/src/datasets/datasets.md
+++ b/docs/src/datasets/datasets.md
@@ -45,7 +45,7 @@ make_output_domain
 ## Underlying data layouts
 
 ```@docs
-AbstractDataLayout
+AbstractLayout
 OneToOne
 ContiguouslyBinned
 common_support
diff --git a/docs/src/examples/optimizers.md b/docs/src/examples/optimizers.md
index 173abe60..2a74c97e 100644
--- a/docs/src/examples/optimizers.md
+++ b/docs/src/examples/optimizers.md
@@ -10,4 +10,16 @@ DATADIR = length(get(ENV, "CI", "")) > 0 ? @__DIR__() * "/../../ex-datadir" : "/
 spec1_path = joinpath(DATADIR, "s54405.pha")
 data = OGIPDataset(spec1_path) 
 normalize!(data)
+
+mask_energies!(data, 1, 15)
+
+# a plotting utility
+my_plot(data) = plot(
+    data, 
+    xscale = :log10, 
+    yscale = :log10,
+    ylims = (1e-3, 1.3)
+)
+
+my_plot(data)
 ```
\ No newline at end of file
diff --git a/docs/src/examples/sherpa-example.md b/docs/src/examples/sherpa-example.md
index ace15e32..ddf80608 100644
--- a/docs/src/examples/sherpa-example.md
+++ b/docs/src/examples/sherpa-example.md
@@ -22,7 +22,7 @@ y_noisy = y .+ (0.2 * randn(length(y)))
 scatter(x, y_noisy)
 ```
 
-To make this into a fittable dataset, we observe that our layout is injective (i.e. `length(x) == length(y)`). This is subtly different from how the majority of spectral models are implemented, which usually assume some kind of binning (`length(x) == length(y) + 1`). Fortunately, SpectralFitting.jl can track this for us, and do various conversion to make the models work correctly for the data. We need only tell the package what our [`AbstractDataLayout`](@ref) is:
+To make this into a fittable dataset, we observe that our layout is injective (i.e. `length(x) == length(y)`). This is subtly different from how the majority of spectral models are implemented, which usually assume some kind of binning (`length(x) == length(y) + 1`). Fortunately, SpectralFitting.jl can track this for us, and do various conversion to make the models work correctly for the data. We need only tell the package what our [`AbstractLayout`](@ref) is:
 
 ```@example sherpa
 data = InjectiveData(x, y_noisy; name = "example")
diff --git a/docs/src/models/using-models.md b/docs/src/models/using-models.md
index 7d6037e4..d1532d51 100644
--- a/docs/src/models/using-models.md
+++ b/docs/src/models/using-models.md
@@ -37,7 +37,7 @@ invokemodel(domain, model)
     ```julia
     length(flux) == length(energy) - 1
     ```
-    Models need not be defined as such, however. See [`AbstractDataLayout`](@ref) for more.
+    Models need not be defined as such, however. See [`AbstractLayout`](@ref) for more.
 
 Models can be composed together following the [Model algebra](@ref). That means to expressive a photoelectric absorption component acting on the power law we can write
 
diff --git a/src/datasets/binning.jl b/src/datasets/binning.jl
index 99396ebf..4e162ec8 100644
--- a/src/datasets/binning.jl
+++ b/src/datasets/binning.jl
@@ -40,21 +40,21 @@ function construct_objective_cache(
     construct_objective_cache(preferred_support(model), T, model, domain)
 end
 function construct_objective_cache(
-    layout::AbstractDataLayout,
+    layout::AbstractLayout,
     model::AbstractSpectralModel,
     domain::AbstractVector{T},
 ) where {T}
     construct_objective_cache(layout, T, model, domain)
 end
 function construct_objective_cache(
-    layout::AbstractDataLayout,
+    layout::AbstractLayout,
     T::Type,
     model::AbstractSpectralModel,
     domain,
 )
     construct_objective_cache(layout, T, length(domain), objective_cache_count(model))
 end
-function construct_objective_cache(layout::AbstractDataLayout, T::Type, N::Int, size::Int)
+function construct_objective_cache(layout::AbstractLayout, T::Type, N::Int, size::Int)
     n = if layout isa OneToOne
         N
     elseif layout isa ContiguouslyBinned
diff --git a/src/datasets/datasets.jl b/src/datasets/datasets.jl
index 929cde69..974a1b85 100644
--- a/src/datasets/datasets.jl
+++ b/src/datasets/datasets.jl
@@ -8,7 +8,7 @@ Fitting data is considered to have an *objective* and a *domain*. As
 the domain may be, for example, energy bins (high and low), or
 fourier frequencies (single value), the purpose of this abstraction
 is to provide some facility for translating between these representations
-for the models to fit with. This is done by checking that the [`AbstractDataLayout`](@ref)
+for the models to fit with. This is done by checking that the [`AbstractLayout`](@ref)
 of the model and data are compatible, or at least have compatible translations.
 
 Must implement a minimal set of accessor methods. These are paired with
@@ -17,7 +17,7 @@ Must implement a minimal set of accessor methods. These are paired with
 into the translation. Usage of these functions should be sparse in the interest of
 performance.
 
-The arrays returned by the `make_*` functions must correspond to the [`AbstractDataLayout`](@ref)
+The arrays returned by the `make_*` functions must correspond to the [`AbstractLayout`](@ref)
 specified by the caller.
 
 - [`make_objective_variance`](@ref)
@@ -46,10 +46,10 @@ end
 
 
 """
-    make_objective(layout::AbstractDataLayout, dataset::AbstractDataset)
+    make_objective(layout::AbstractLayout, dataset::AbstractDataset)
 
 Returns the array used as the target for model fitting. The array must
-correspond to the data [`AbstractDataLayout`](@ref) specified by the `layout`
+correspond to the data [`AbstractLayout`](@ref) specified by the `layout`
 parameter.
 
 In as far as it can be guarunteed, the memory in the returned array will not be
@@ -57,36 +57,36 @@ mutated by any fitting procedures.
 
 Domain for this objective should be returned by [`make_model_domain`](@ref).
 """
-make_objective(layout::AbstractDataLayout, dataset::AbstractDataset) =
+make_objective(layout::AbstractLayout, dataset::AbstractDataset) =
     error("Layout $(layout) is not implemented for $(typeof(dataset))")
 
 """
-    make_objective_variance(layout::AbstractDataLayout, dataset::AbstractDataset)
+    make_objective_variance(layout::AbstractLayout, dataset::AbstractDataset)
 
 Make the variance vector associated with each objective point.
 """
-make_objective_variance(layout::AbstractDataLayout, dataset::AbstractDataset) =
+make_objective_variance(layout::AbstractLayout, dataset::AbstractDataset) =
     error("Layout $(layout) is not implemented for $(typeof(dataset))")
 
 """
-    make_model_domain(layout::AbstractDataLayout, dataset::AbstractDataset)
+    make_model_domain(layout::AbstractLayout, dataset::AbstractDataset)
 
 Returns the array used as the domain for the modelling. This is paired with
 [`make_domain_variance`](@ref)
 """
-make_model_domain(layout::AbstractDataLayout, dataset::AbstractDataset) =
+make_model_domain(layout::AbstractLayout, dataset::AbstractDataset) =
     error("Layout $(layout) is not implemented for $(typeof(dataset))")
 
 """
-    make_domain_variance(layout::AbstractDataLayout, dataset::AbstractDataset) 
+    make_domain_variance(layout::AbstractLayout, dataset::AbstractDataset) 
 
 Make the variance vector associated with the domain.
 """
-make_domain_variance(layout::AbstractDataLayout, dataset::AbstractDataset) =
+make_domain_variance(layout::AbstractLayout, dataset::AbstractDataset) =
     error("Layout $(layout) is not implemented for $(typeof(dataset))")
 
 """
-    make_output_domain(layout::AbstractDataLayout, dataset::AbstractDataset)
+    make_output_domain(layout::AbstractLayout, dataset::AbstractDataset)
 
 Returns the array used as the output domain. That is, in cases where the model
 input and output map to different domains, the input domain is said to be the
@@ -95,11 +95,11 @@ model domain, the input domain is said to be the model domain.
 The distinction is mainly used for the purposes of simulating data and for
 visualising data.
 """
-make_output_domain(layout::AbstractDataLayout, dataset::AbstractDataset) =
+make_output_domain(layout::AbstractLayout, dataset::AbstractDataset) =
     error("Layout $(layout) is not implemented for $(typeof(dataset))")
 
 """
-    objective_transformer(layout::AbstractDataLayout, dataset::AbstractDataset)
+    objective_transformer(layout::AbstractLayout, dataset::AbstractDataset)
 
 Used to transform the output of the model onto the output domain. For spectral
 fitting, this is the method used to do response folding and bin masking.
@@ -117,7 +117,7 @@ function _DEFAULT_TRANSFORMER()
 end
 ```
 """
-function objective_transformer(layout::AbstractDataLayout, dataset::AbstractDataset)
+function objective_transformer(layout::AbstractLayout, dataset::AbstractDataset)
     @warn "Using default objective transformer!"
     _DEFAULT_TRANSFORMER()
 end
diff --git a/src/datasets/injectivedata.jl b/src/datasets/injectivedata.jl
index f60cde25..3ab979d8 100644
--- a/src/datasets/injectivedata.jl
+++ b/src/datasets/injectivedata.jl
@@ -35,13 +35,10 @@ make_objective(::ContiguouslyBinned, dataset::InjectiveData) =
 make_model_domain(::OneToOne, dataset::InjectiveData) = dataset.domain
 make_objective(::OneToOne, dataset::InjectiveData) = dataset.codomain[dataset.data_mask]
 
-make_output_domain(layout::AbstractDataLayout, dataset::InjectiveData) =
+make_output_domain(layout::AbstractLayout, dataset::InjectiveData) =
     make_model_domain(layout, dataset)
 
-function make_objective_variance(
-    ::AbstractDataLayout,
-    dataset::InjectiveData{V},
-)::V where {V}
+function make_objective_variance(::AbstractLayout, dataset::InjectiveData{V})::V where {V}
     if !isnothing(dataset.codomain_variance)
         dataset.codomain_variance[dataset.data_mask]
     else
@@ -50,7 +47,7 @@ function make_objective_variance(
     end
 end
 
-function objective_transformer(::AbstractDataLayout, dataset::InjectiveData)
+function objective_transformer(::AbstractLayout, dataset::InjectiveData)
     function _transformer!!(domain, objective)
         @views objective[dataset.data_mask]
     end
diff --git a/src/datasets/spectraldata.jl b/src/datasets/spectraldata.jl
index 38426bcd..70f58b80 100644
--- a/src/datasets/spectraldata.jl
+++ b/src/datasets/spectraldata.jl
@@ -168,12 +168,12 @@ function check_units_warning(units)
     end
 end
 
-function make_objective(layout::AbstractDataLayout, dataset::SpectralData)
+function make_objective(layout::AbstractLayout, dataset::SpectralData)
     check_units_warning(dataset.spectrum.units)
     make_objective(layout, dataset.spectrum)[dataset.data_mask]
 end
 
-function make_objective_variance(layout::AbstractDataLayout, dataset::SpectralData)
+function make_objective_variance(layout::AbstractLayout, dataset::SpectralData)
     check_units_warning(dataset.spectrum.units)
     make_objective_variance(layout, dataset.spectrum)[dataset.data_mask]
 end
@@ -427,27 +427,23 @@ macro _forward_SpectralData_api(args)
     quote
         SpectralFitting.supports(::ContiguouslyBinned, t::Type{<:$(T)}) = true
         SpectralFitting.make_output_domain(
-            layout::SpectralFitting.AbstractDataLayout,
+            layout::SpectralFitting.AbstractLayout,
             t::$(T),
         ) = SpectralFitting.make_output_domain(layout, getfield(t, $(field)))
-        SpectralFitting.make_model_domain(
-            layout::SpectralFitting.AbstractDataLayout,
-            t::$(T),
-        ) = SpectralFitting.make_model_domain(layout, getfield(t, $(field)))
+        SpectralFitting.make_model_domain(layout::SpectralFitting.AbstractLayout, t::$(T)) =
+            SpectralFitting.make_model_domain(layout, getfield(t, $(field)))
         SpectralFitting.make_domain_variance(
-            layout::SpectralFitting.AbstractDataLayout,
+            layout::SpectralFitting.AbstractLayout,
             t::$(T),
         ) = SpectralFitting.make_domain_variance(layout, getfield(t, $(field)))
-        SpectralFitting.make_objective(
-            layout::SpectralFitting.AbstractDataLayout,
-            t::$(T),
-        ) = SpectralFitting.make_objective(layout, getfield(t, $(field)))
+        SpectralFitting.make_objective(layout::SpectralFitting.AbstractLayout, t::$(T)) =
+            SpectralFitting.make_objective(layout, getfield(t, $(field)))
         SpectralFitting.make_objective_variance(
-            layout::SpectralFitting.AbstractDataLayout,
+            layout::SpectralFitting.AbstractLayout,
             t::$(T),
         ) = SpectralFitting.make_objective_variance(layout, getfield(t, $(field)))
         SpectralFitting.objective_transformer(
-            layout::SpectralFitting.AbstractDataLayout,
+            layout::SpectralFitting.AbstractLayout,
             t::$(T),
         ) = SpectralFitting.objective_transformer(layout, getfield(t, $(field)))
         SpectralFitting.regroup!(t::$(T), args...; kwargs...) =
diff --git a/src/fitting/cache.jl b/src/fitting/cache.jl
index 787099c3..f9f3293a 100644
--- a/src/fitting/cache.jl
+++ b/src/fitting/cache.jl
@@ -12,7 +12,7 @@ struct SpectralCache{M,O,T,K,P,TransformerType} <: AbstractFittingCache
     parameter_cache::P
     transformer!!::TransformerType
     function SpectralCache(
-        layout::AbstractDataLayout,
+        layout::AbstractLayout,
         model::M,
         domain,
         objective,
@@ -130,9 +130,9 @@ function _build_mapping_length(f, itt::Tuple)
     values, mapping
 end
 
-_build_objective_mapping(layout::AbstractDataLayout, dataset::FittableMultiDataset) =
+_build_objective_mapping(layout::AbstractLayout, dataset::FittableMultiDataset) =
     _build_mapping_length(i -> make_objective(layout, i), dataset.d)
-_build_domain_mapping(layout::AbstractDataLayout, dataset::FittableMultiDataset) =
+_build_domain_mapping(layout::AbstractLayout, dataset::FittableMultiDataset) =
     _build_mapping_length(i -> make_model_domain(layout, i), dataset.d)
-_build_output_domain_mapping(layout::AbstractDataLayout, dataset::FittableMultiDataset) =
+_build_output_domain_mapping(layout::AbstractLayout, dataset::FittableMultiDataset) =
     _build_mapping_length(i -> make_output_domain(layout, i), dataset.d)
diff --git a/src/support.jl b/src/support.jl
index 404c4716..95a004f7 100644
--- a/src/support.jl
+++ b/src/support.jl
@@ -6,7 +6,7 @@
 end
 
 """
-    abstract type AbstractDataLayout end
+    abstract type AbstractLayout end
 
 The data layout primarily concerns the relationship between the objective and
 the domain. It is used to work out whether a model and a dataset are fittable,
@@ -20,10 +20,10 @@ The following methods may be used to interrogate support:
 The following method is also used to define the support of a model or dataset:
 - [`supports`](@ref)
 """
-abstract type AbstractDataLayout end
+abstract type AbstractLayout end
 
 """
-    struct OneToOne <: AbstractDataLayout end
+    struct OneToOne <: AbstractLayout end
 
 Indicates there is a one-to-one (injective) correspondence between each input
 value and each output value. That is to say
@@ -31,10 +31,10 @@ value and each output value. That is to say
 length(objective) == length(domain)
 ```
 """
-struct OneToOne <: AbstractDataLayout end
+struct OneToOne <: AbstractLayout end
 
 """
-    struct ContiguouslyBinned <: AbstractDataLayout end
+    struct ContiguouslyBinned <: AbstractLayout end
 
 Contiguously binned data layout means that the domain describes high and low
 bins, with the objective being the value in that bin. This means
@@ -47,14 +47,14 @@ Note that the _contiguous_ qualifer is to mean there is no gaps in the bins, and
 ```
 
 """
-struct ContiguouslyBinned <: AbstractDataLayout end
+struct ContiguouslyBinned <: AbstractLayout end
 
 const DEFAULT_SUPPORT_ORDERING = (ContiguouslyBinned(), OneToOne())
 
 """
     preferred_support(x)
 
-Get the preffered [`AbstractDataLayout`](@ref) of `x`. If multiple supports are available, 
+Get the preffered [`AbstractLayout`](@ref) of `x`. If multiple supports are available, 
 the `DEFAULT_SUPPORT_ORDERING` is followed:
 
 ```
@@ -73,7 +73,7 @@ end
 """
     common_support(x, y)
 
-Find the common [`AbstractDataLayout`](@ref) of `x` and `y`, following the ordering of
+Find the common [`AbstractLayout`](@ref) of `x` and `y`, following the ordering of
 [`preferred_support`](@ref).
 """
 function common_support(x, y)
@@ -110,19 +110,19 @@ end
 common_support(args::Vararg) = reduce(_support_reducer, args)
 
 """
-    supports(layout::AbstractDataLayout, x::Type)::Bool
+    supports(layout::AbstractLayout, x::Type)::Bool
 
 Used to define whether a given type has support for a specific
-[`AbstractDataLayout`](@ref). This method should be implemented to express new
+[`AbstractLayout`](@ref). This method should be implemented to express new
 support, not the query method.
 
 To query, there is
 
-    support(layout::AbstractDataLayout, x)::Bool
+    support(layout::AbstractLayout, x)::Bool
 """
-supports(layout::AbstractDataLayout, x) = supports(layout, typeof(x))
+supports(layout::AbstractLayout, x) = supports(layout, typeof(x))
 supports(::ContiguouslyBinned, T::Type) = false
 supports(::OneToOne, T::Type) = false
 
 export OneToOne,
-    ContiguouslyBinned, AbstractDataLayout, supports, preferred_support, common_support
+    ContiguouslyBinned, AbstractLayout, supports, preferred_support, common_support

From c2d4c62c0352dbf01f4e3130313bcf87dbb7243a Mon Sep 17 00:00:00 2001
From: fjebaker <fergusbkr@gmail.com>
Date: Sat, 25 May 2024 23:45:18 +0100
Subject: [PATCH 3/6] feat!: better support API

Changed the function `supports` to now return a list of the supported
layouts instead of having to define a new dispatch for each one.
---
 src/abstract-models.jl                  |  2 +-
 src/datasets/injectivedata.jl           |  3 +--
 src/datasets/spectraldata.jl            |  4 ++--
 src/datasets/spectrum.jl                |  2 +-
 src/simulate.jl                         |  2 +-
 src/support.jl                          | 21 ++++++++++++++-------
 test/fitting/test-fit-simple-dataset.jl |  3 +--
 7 files changed, 21 insertions(+), 16 deletions(-)

diff --git a/src/abstract-models.jl b/src/abstract-models.jl
index d562dc06..9f5ce043 100644
--- a/src/abstract-models.jl
+++ b/src/abstract-models.jl
@@ -96,7 +96,7 @@ Model reflection is supported by the following functions. These are intended for
 The parametric type parameter `T` is the number type of the model and `K` defines the [`AbstractSpectralModelKind`](@ref).
 """
 abstract type AbstractSpectralModel{T,K<:AbstractSpectralModelKind} end
-supports(::ContiguouslyBinned, ::Type{<:AbstractSpectralModel}) = true
+supports(::Type{<:AbstractSpectralModel}) = (ContiguouslyBinned(),)
 
 """
     numbertype(::AbstractSpectralModel)
diff --git a/src/datasets/injectivedata.jl b/src/datasets/injectivedata.jl
index 3ab979d8..70215bf5 100644
--- a/src/datasets/injectivedata.jl
+++ b/src/datasets/injectivedata.jl
@@ -19,8 +19,7 @@ function InjectiveData(
     InjectiveData(domain, codomain, domain_variance, codomain_variance, name, data_mask)
 end
 
-supports(::ContiguouslyBinned, ::Type{<:InjectiveData}) = true
-supports(::OneToOne, ::Type{<:InjectiveData}) = true
+supports(::Type{<:InjectiveData}) = (ContiguouslyBinned(), OneToOne())
 
 function make_model_domain(::ContiguouslyBinned, dataset::InjectiveData)
     # need to expand the domain by one
diff --git a/src/datasets/spectraldata.jl b/src/datasets/spectraldata.jl
index 70f58b80..84dc4c20 100644
--- a/src/datasets/spectraldata.jl
+++ b/src/datasets/spectraldata.jl
@@ -160,7 +160,7 @@ function SpectralData(
     return data
 end
 
-supports(::ContiguouslyBinned, ::Type{<:SpectralData}) = true
+supports(::Type{<:SpectralData}) = (ContiguouslyBinned(),)
 
 function check_units_warning(units)
     if units != u"counts / (s * keV)"
@@ -425,7 +425,7 @@ macro _forward_SpectralData_api(args)
     end
     T, field = args.args
     quote
-        SpectralFitting.supports(::ContiguouslyBinned, t::Type{<:$(T)}) = true
+        SpectralFitting.supports(t::Type{<:$(T)}) = (ContiguouslyBinned(),)
         SpectralFitting.make_output_domain(
             layout::SpectralFitting.AbstractLayout,
             t::$(T),
diff --git a/src/datasets/spectrum.jl b/src/datasets/spectrum.jl
index 9100e424..00f6ce03 100644
--- a/src/datasets/spectrum.jl
+++ b/src/datasets/spectrum.jl
@@ -28,7 +28,7 @@ function normalize!(spectrum::Spectrum)
     spectrum
 end
 
-supports(::ContiguouslyBinned, ::Type{<:Spectrum}) = true
+supports(::Type{<:Spectrum}) = (ContiguouslyBinned(),)
 
 function make_objective(::ContiguouslyBinned, dataset::Spectrum)
     check_units_warning(dataset.units)
diff --git a/src/simulate.jl b/src/simulate.jl
index b1cb4b62..8e71e15c 100644
--- a/src/simulate.jl
+++ b/src/simulate.jl
@@ -9,7 +9,7 @@ mutable struct SimulatedSpectrum{T,F} <: AbstractDataset
     seed::Int
 end
 
-supports(::ContiguouslyBinned, ::Type{<:SimulatedSpectrum}) = true
+supports(::Type{<:SimulatedSpectrum}) = (ContiguouslyBinned(),)
 
 function make_objective(::ContiguouslyBinned, dataset::SimulatedSpectrum)
     check_units_warning(dataset.units)
diff --git a/src/support.jl b/src/support.jl
index 95a004f7..753623dd 100644
--- a/src/support.jl
+++ b/src/support.jl
@@ -110,19 +110,26 @@ end
 common_support(args::Vararg) = reduce(_support_reducer, args)
 
 """
-    supports(layout::AbstractLayout, x::Type)::Bool
+    supports(x::Type)
 
 Used to define whether a given type has support for a specific
-[`AbstractLayout`](@ref). This method should be implemented to express new
-support, not the query method.
+[`AbstractLayout`](@ref). Should return a tuple of the supported layouts. This
+method should be implemented to express new support, not the query method. 
 
 To query, there is
 
-    support(layout::AbstractLayout, x)::Bool
+    supports(layout::AbstractLayout, x)::Bool
+
+# Example
+
+```julia
+supports(::Type{typeof(x)}) = (OneToOne(),)
+@assert supports(ContiguouslyBinned(), x) == false
+```
 """
-supports(layout::AbstractLayout, x) = supports(layout, typeof(x))
-supports(::ContiguouslyBinned, T::Type) = false
-supports(::OneToOne, T::Type) = false
+supports(layout::AbstractLayout, x)::Bool = layout in supports(x)
+supports(::T) where {T} = supports(T)
+supports(::Type) = ()
 
 export OneToOne,
     ContiguouslyBinned, AbstractLayout, supports, preferred_support, common_support
diff --git a/test/fitting/test-fit-simple-dataset.jl b/test/fitting/test-fit-simple-dataset.jl
index 6ed469a5..432063ea 100644
--- a/test/fitting/test-fit-simple-dataset.jl
+++ b/test/fitting/test-fit-simple-dataset.jl
@@ -10,8 +10,7 @@ end
 function SpectralFitting.invoke!(out, x, model::DummySimpleLinear)
     @. out = x + (model.c / model.K)
 end
-SpectralFitting.supports(::ContiguouslyBinned, ::Type{<:DummySimpleLinear}) = false
-SpectralFitting.supports(::OneToOne, ::Type{<:DummySimpleLinear}) = true
+SpectralFitting.supports(::Type{<:DummySimpleLinear}) = (OneToOne(),)
 
 x = collect(range(1.0, 10.0, 101))
 y = @. 13.12 * x + 102

From 42ed291e4ecc4e54b08447bb34b375141788b254 Mon Sep 17 00:00:00 2001
From: fjebaker <fergusbkr@gmail.com>
Date: Sat, 25 May 2024 23:53:41 +0100
Subject: [PATCH 4/6] feat: add support_units for adding unit information to
 layouts

---
 src/support.jl | 28 ++++++++++++++++++++++++----
 1 file changed, 24 insertions(+), 4 deletions(-)

diff --git a/src/support.jl b/src/support.jl
index 753623dd..2c2aac74 100644
--- a/src/support.jl
+++ b/src/support.jl
@@ -19,9 +19,21 @@ The following methods may be used to interrogate support:
 
 The following method is also used to define the support of a model or dataset:
 - [`supports`](@ref)
+
+For cases where unit information needs to be propagated, an `AbstractLayout` can also be used to ensure the units are compatible. To query the units of a layout, use
+- [`support_units`](@ref)
 """
 abstract type AbstractLayout end
 
+"""
+    support_units(x)
+
+Return the units of a particular layout. If this method returns `nothing`,
+assume the layout does not care about the units and handle that information
+appropriately (throw an error or set defaults).
+"""
+support_units(::T) where {T<:AbstractLayout} = nothing
+
 """
     struct OneToOne <: AbstractLayout end
 
@@ -31,7 +43,11 @@ value and each output value. That is to say
 length(objective) == length(domain)
 ```
 """
-struct OneToOne <: AbstractLayout end
+struct OneToOne{U} <: AbstractLayout
+    units::U
+end
+OneToOne() = OneToOne(nothing)
+support_units(l::OneToOne) = l.units
 
 """
     struct ContiguouslyBinned <: AbstractLayout end
@@ -47,7 +63,11 @@ Note that the _contiguous_ qualifer is to mean there is no gaps in the bins, and
 ```
 
 """
-struct ContiguouslyBinned <: AbstractLayout end
+struct ContiguouslyBinned{U} <: AbstractLayout
+    units::U
+end
+ContiguouslyBinned() = ContiguouslyBinned(nothing)
+support_units(l::ContiguouslyBinned) = l.units
 
 const DEFAULT_SUPPORT_ORDERING = (ContiguouslyBinned(), OneToOne())
 
@@ -131,5 +151,5 @@ supports(layout::AbstractLayout, x)::Bool = layout in supports(x)
 supports(::T) where {T} = supports(T)
 supports(::Type) = ()
 
-export OneToOne,
-    ContiguouslyBinned, AbstractLayout, supports, preferred_support, common_support
+export AbstractLayout,
+    common_support, ContiguouslyBinned, OneToOne, preferred_support, supports, support_units

From ddcf631461bbb2345a126814f70a208813dc7008 Mon Sep 17 00:00:00 2001
From: fjebaker <fergusbkr@gmail.com>
Date: Sun, 26 May 2024 01:36:27 +0100
Subject: [PATCH 5/6] feat: add units API to AbstractDataset

Add a specification of how the abstract dataset should indicate
preferred units per fitting statistic, so that different statistics can
be easily fit in different regimes.
---
 src/SpectralFitting.jl       |  2 ++
 src/datasets/datasets.jl     | 52 +++++++++++++++++++++++++++++++++---
 src/datasets/spectraldata.jl |  6 +++++
 src/fitting/config.jl        |  2 +-
 src/fitting/statistics.jl    |  2 --
 src/simulate.jl              |  2 --
 src/support.jl               | 12 ++++++++-
 test/datasets/test-units.jl  | 13 +++++++++
 test/runtests.jl             | 15 ++++++-----
 9 files changed, 89 insertions(+), 17 deletions(-)
 create mode 100644 test/datasets/test-units.jl

diff --git a/src/SpectralFitting.jl b/src/SpectralFitting.jl
index 43e15780..c795457a 100644
--- a/src/SpectralFitting.jl
+++ b/src/SpectralFitting.jl
@@ -36,6 +36,8 @@ abstract type AbstractMission end
 struct NoMission <: AbstractMission end
 
 abstract type AbstractStatistic end
+struct ChiSquared <: AbstractStatistic end
+struct Cash <: AbstractStatistic end
 
 # unitful units
 include("units.jl")
diff --git a/src/datasets/datasets.jl b/src/datasets/datasets.jl
index 974a1b85..584cc86a 100644
--- a/src/datasets/datasets.jl
+++ b/src/datasets/datasets.jl
@@ -30,6 +30,15 @@ Additionally there is an objective transformer that transforms the output of the
 model onto the `output` domain:
 
 - [`objective_transformer`](@ref)
+
+Finally, to make all of the fitting for different statistical regimes work
+efficiently, datasets should inform which units are preferred to fit. They may
+also give the error statistics they prefer, and a label name primarily used to
+disambiguate:
+
+- [`preferred_units`](@ref)
+- [`error_statistic`](@ref)
+- [`make_label`](@ref)
 """
 abstract type AbstractDataset end
 
@@ -132,23 +141,58 @@ function _DEFAULT_TRANSFORMER()
     _transformer!!
 end
 
-make_label(d::AbstractDataset) = "$(Base.typename(typeof(d)).name)"
+"""
+    preferred_units(::Type{<:AbstractDataset}, s::AbstractStatistic)
+    preferred_units(x, s::AbstractStatistic)
+
+Get the preferred units that a given dataset would use to fit the
+[`AbstractStatistic`](@ref) in. For example, for [`ChiSquared`](@ref), the units
+of the model may be a rate, however for [`Cash`](@ref) the preferred units might
+be counts.
+
+Returning `nothing` from this function implies there is no unit preference.
+
+If undefined for a derived type, returns `nothing`.
+
+See also [`support_units`](@ref).
+"""
+preferred_units(::T, s::AbstractStatistic) where {T<:AbstractDataset} =
+    preferred_units(T, s)
+preferred_units(::Type{<:AbstractDataset}, s::AbstractStatistic) = nothing
+
+"""
+    error_statistic(::AbstractDataset)
 
+Should return an [`ErrorStatistics`](@ref) describing which error statistic this
+data uses.
+
+If undefined for a derived type, returns `ErrorStatistics.Unknown`.
+"""
 error_statistic(::AbstractDataset) = ErrorStatistics.Unknown
 
+"""
+    make_label(d::AbstractDataset)
+
+Return a string that gives a descriptive label for this dataset.
+"""
+make_label(d::AbstractDataset) = "$(Base.typename(typeof(d)).name)"
+
 """
 Must support the same API, but may also have some query methods for specific internals.
 """
 abstract type AbstractMultiDataset <: AbstractDataset end
 
 export AbstractDataset,
-    make_model_domain,
-    make_output_domain,
+    error_statistic,
     make_domain_variance,
+    make_label,
+    make_model_domain,
     make_objective,
     make_objective_variance,
+    make_output_domain,
     normalize!,
-    objective_transformer
+    objective_transformer,
+    preferred_units
 
 include("spectrum.jl")
 include("response.jl")
diff --git a/src/datasets/spectraldata.jl b/src/datasets/spectraldata.jl
index 84dc4c20..8aa126b1 100644
--- a/src/datasets/spectraldata.jl
+++ b/src/datasets/spectraldata.jl
@@ -329,6 +329,10 @@ objective_units(data::SpectralData) = data.spectrum.units
 
 error_statistic(data::SpectralData) = error_statistic(data.spectrum)
 
+preferred_units(::Type{<:SpectralData}, ::AbstractStatistic) = nothing
+preferred_units(::Type{<:SpectralData}, ::ChiSquared) = u"counts / (s * keV)"
+preferred_units(::Type{<:SpectralData}, ::Cash) = u"counts"
+
 # internal methods
 
 function rebin_if_different_domains!(output, data_domain, model_domain, input)
@@ -426,6 +430,8 @@ macro _forward_SpectralData_api(args)
     T, field = args.args
     quote
         SpectralFitting.supports(t::Type{<:$(T)}) = (ContiguouslyBinned(),)
+        SpectralFitting.preferred_units(t::Type{<:$(T)}, u::AbstractStatistic) =
+            SpectralFitting.preferred_units(SpectralData, u)
         SpectralFitting.make_output_domain(
             layout::SpectralFitting.AbstractLayout,
             t::$(T),
diff --git a/src/fitting/config.jl b/src/fitting/config.jl
index 566a621f..3dd82bcd 100644
--- a/src/fitting/config.jl
+++ b/src/fitting/config.jl
@@ -41,7 +41,7 @@ function make_single_config(prob::FittingProblem, stat::AbstractStatistic)
     model = prob.model.m[1]
     dataset = prob.data.d[1]
 
-    layout = common_support(model, dataset)
+    layout = with_units(common_support(model, dataset), preferred_units(dataset, stat))
     model_domain = make_model_domain(layout, dataset)
     output_domain = make_output_domain(layout, dataset)
     objective = make_objective(layout, dataset)
diff --git a/src/fitting/statistics.jl b/src/fitting/statistics.jl
index 10d56770..654fbf5b 100644
--- a/src/fitting/statistics.jl
+++ b/src/fitting/statistics.jl
@@ -2,7 +2,6 @@ function measure(stat::AbstractStatistic, result::FittingResult, args...; kwargs
     measure(stat, result[1], args...; kwargs...)
 end
 
-struct ChiSquared <: AbstractStatistic end
 function measure(::ChiSquared, y, ŷ, σ²)
     sum(@.((y - ŷ)^2 / σ²))
 end
@@ -22,7 +21,6 @@ function _f_wrap_objective(stat::ChiSquared, config::FittingConfig)
     end
 end
 
-struct Cash <: AbstractStatistic end
 function measure(::Cash, S, ŷ)
     2 * sum(@.(ŷ - S + S * (log(S) - log(ŷ))))
 end
diff --git a/src/simulate.jl b/src/simulate.jl
index 8e71e15c..90fee63b 100644
--- a/src/simulate.jl
+++ b/src/simulate.jl
@@ -12,12 +12,10 @@ end
 supports(::Type{<:SimulatedSpectrum}) = (ContiguouslyBinned(),)
 
 function make_objective(::ContiguouslyBinned, dataset::SimulatedSpectrum)
-    check_units_warning(dataset.units)
     dataset.data
 end
 
 function make_objective_variance(::ContiguouslyBinned, dataset::SimulatedSpectrum)
-    check_units_warning(dataset.units)
     dataset.variance
 end
 
diff --git a/src/support.jl b/src/support.jl
index 2c2aac74..d9281996 100644
--- a/src/support.jl
+++ b/src/support.jl
@@ -34,6 +34,14 @@ appropriately (throw an error or set defaults).
 """
 support_units(::T) where {T<:AbstractLayout} = nothing
 
+"""
+    with_units(::AbstractLayout, units)
+
+Remake the [`AbstractLayout`](@ref) with the desired units. This may be a no-op
+if the layout does not care about units, see [`support_units`](@ref).
+"""
+with_units(layout::AbstractLayout, units) = layout
+
 """
     struct OneToOne <: AbstractLayout end
 
@@ -48,6 +56,7 @@ struct OneToOne{U} <: AbstractLayout
 end
 OneToOne() = OneToOne(nothing)
 support_units(l::OneToOne) = l.units
+with_units(::OneToOne, units) = OneToOne(units)
 
 """
     struct ContiguouslyBinned <: AbstractLayout end
@@ -68,13 +77,14 @@ struct ContiguouslyBinned{U} <: AbstractLayout
 end
 ContiguouslyBinned() = ContiguouslyBinned(nothing)
 support_units(l::ContiguouslyBinned) = l.units
+with_units(::ContiguouslyBinned, units) = ContiguouslyBinned(units)
 
 const DEFAULT_SUPPORT_ORDERING = (ContiguouslyBinned(), OneToOne())
 
 """
     preferred_support(x)
 
-Get the preffered [`AbstractLayout`](@ref) of `x`. If multiple supports are available, 
+Get the preferred [`AbstractLayout`](@ref) of `x`. If multiple supports are available, 
 the `DEFAULT_SUPPORT_ORDERING` is followed:
 
 ```
diff --git a/test/datasets/test-units.jl b/test/datasets/test-units.jl
new file mode 100644
index 00000000..765d10dd
--- /dev/null
+++ b/test/datasets/test-units.jl
@@ -0,0 +1,13 @@
+using SpectralFitting, Test
+
+include("../dummies.jl")
+
+data = make_dummy_dataset(collect(range(1.0, 0.0, 10)), collect(range(0, 15.0, 11)))
+
+# this really needs to be compile time known
+@inferred preferred_units(data, ChiSquared())
+@inferred preferred_units(typeof(data), ChiSquared())
+
+set_units!(data, u"counts / (s * keV)")
+
+make_objective(SpectralFitting.ContiguouslyBinned(preferred_units(data, Cash())), data)
diff --git a/test/runtests.jl b/test/runtests.jl
index 857a288c..f3c189eb 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -68,15 +68,16 @@ end
     @testset "printing" begin
         include("io/test-printing.jl")
     end
-    if has_test_dir
-        @testset "datasets" begin
+    @testset "datasets" begin
+        if has_test_dir
             include("datasets/test-ogip.jl")
-            include("datasets/test-grouping.jl")
-            include("datasets/test-binning.jl")
-            include("datasets/test-datasets.jl")
+        else
+            @warn "Skipping OGIP dataset tests."
         end
-    else
-        @warn "Skipping dataset tests."
+        include("datasets/test-grouping.jl")
+        include("datasets/test-units.jl")
+        include("datasets/test-binning.jl")
+        include("datasets/test-datasets.jl")
     end
     include("io/test-remote-pathname-compression.jl")
 end

From 5577b718768abe71651960c02c58e5afd7cc8cc3 Mon Sep 17 00:00:00 2001
From: fjebaker <fergusbkr@gmail.com>
Date: Sun, 26 May 2024 02:36:38 +0100
Subject: [PATCH 6/6] feat: units in fitting

Started adding the code that tracks that the model, data, and fit
statistic are used to change the units used when fitting. For example,
when using Cash stat, we want to fit in counts, but for ChiSquared we
want to fit in counts / (s * keV^-1), or similar.

There is a looming issue here related to type instability with the
units, since the units are basically compile time structs, but the way
we store them in e.g. Spectrum is as a union so they may easily be
modified.

This means we can't rely on the dataset units for codegen.

deprecated: normalize!
---
 src/datasets/spectraldata.jl | 127 ++++++++++++++++++++++-------------
 src/datasets/spectrum.jl     |   2 -
 src/fitting/config.jl        |   7 +-
 src/fitting/statistics.jl    |  18 ++---
 src/simulate.jl              |   2 +-
 5 files changed, 90 insertions(+), 66 deletions(-)

diff --git a/src/datasets/spectraldata.jl b/src/datasets/spectraldata.jl
index 8aa126b1..0eac96bf 100644
--- a/src/datasets/spectraldata.jl
+++ b/src/datasets/spectraldata.jl
@@ -108,24 +108,6 @@ end
 
 function SpectralData(paths::SpectralDataPaths; kwargs...)
     spec, resp, back, anc = _read_all_ogip(paths; kwargs...)
-
-    # convert everything to rates
-    if spec.units == u"counts"
-        spec.units = u"counts / s"
-        @. spec.data /= spec.exposure_time
-        if !isnothing(spec.errors)
-            @. spec.errors /= spec.exposure_time
-        end
-    end
-
-    if !isnothing(back) && back.units == u"counts"
-        back.units = u"counts / s"
-        @. back.data /= back.exposure_time
-        if !isnothing(back.errors)
-            @. back.errors /= back.exposure_time
-        end
-    end
-
     SpectralData(spec, resp; background = back, ancillary = anc)
 end
 
@@ -162,20 +144,25 @@ end
 
 supports(::Type{<:SpectralData}) = (ContiguouslyBinned(),)
 
-function check_units_warning(units)
-    if units != u"counts / (s * keV)"
-        @warn "Data is currently still in $(units). Most models fit in rate (count / (s * keV)). Use `normalize!(dataset)` to ensure the dataset is in a standard format."
+function _objective_to_units(dataset::SpectralData, obj, units)
+    adj = if !isnothing(units)
+        alt = copy(obj)
+        adjust_to_units!(dataset, dataset.spectrum, alt, units)
+        alt
+    else
+        obj
     end
+    adj[dataset.data_mask]
 end
 
 function make_objective(layout::AbstractLayout, dataset::SpectralData)
-    check_units_warning(dataset.spectrum.units)
-    make_objective(layout, dataset.spectrum)[dataset.data_mask]
+    obj = make_objective(layout, dataset.spectrum)
+    _objective_to_units(dataset, obj, support_units(layout))
 end
 
 function make_objective_variance(layout::AbstractLayout, dataset::SpectralData)
-    check_units_warning(dataset.spectrum.units)
-    make_objective_variance(layout, dataset.spectrum)[dataset.data_mask]
+    var = make_objective_variance(layout, dataset.spectrum)
+    _objective_to_units(dataset, var, support_units(layout))
 end
 
 make_model_domain(::ContiguouslyBinned, dataset::SpectralData) = dataset.domain
@@ -198,17 +185,28 @@ function restrict_domain!(dataset::SpectralData, condition)
     dataset
 end
 
-function _fold_transformer(T::Type, layout, R, ΔE, E)
+function _fold_transformer(T::Type, exposure_time, layout::AbstractLayout, R, ΔE, E)
     cache = DiffCache(construct_objective_cache(layout, T, length(E), 1))
+    units = support_units(layout)
     function _transformer!!(energy, flux)
         f = rebin_if_different_domains!(get_tmp(cache, flux), E, energy, flux)
         f = R * f
-        @. f = f / ΔE
+        if units === u"counts / (s * keV)"
+            @. f = f / ΔE
+        elseif units === u"counts"
+            @. f = f * exposure_time
+        end
+        f
     end
     function _transformer!!(output, energy, flux)
         f = rebin_if_different_domains!(get_tmp(cache, flux), E, energy, flux)
         mul!(output, R, f)
-        @. output = output / ΔE
+        if units === u"counts / (s * keV)"
+            @. output = output / ΔE
+        elseif units === u"counts"
+            @. output = output * exposure_time
+        end
+        output
     end
     _transformer!!
 end
@@ -225,7 +223,7 @@ function objective_transformer(
     R = R_folded[dataset.data_mask, :]
     ΔE = bin_widths(dataset)
     model_domain = response_energy(dataset.response)
-    _fold_transformer(T, layout, R, ΔE, model_domain)
+    _fold_transformer(T, dataset.spectrum.exposure_time, layout, R, ΔE, model_domain)
 end
 
 
@@ -294,22 +292,8 @@ function regroup!(dataset::SpectralData; min_counts = nothing)
 end
 
 function normalize!(dataset::SpectralData)
-    ΔE = bin_widths(dataset)
-    normalize!(dataset.spectrum)
-    if !(dataset.spectrum.units == u"counts / (s * keV)")
-        @. dataset.spectrum.data /= ΔE
-        @. dataset.spectrum.errors /= ΔE
-        dataset.spectrum.units = u"counts / (s * keV)"
-    end
-    if has_background(dataset)
-        normalize!(dataset.background)
-        if !(dataset.background.units == u"counts / (s * keV)")
-            @. dataset.background.data /= ΔE
-            @. dataset.background.errors /= ΔE
-            dataset.background.units = u"counts / (s * keV)"
-        end
-    end
-    dataset
+    Base.depwarn("`normalize!` is deprecated. Use `set_units!` instead.", :normalize!)
+    set_units!(dataset, u"counts / (s * keV)")
 end
 
 function subtract_background!(dataset::SpectralData)
@@ -333,6 +317,53 @@ preferred_units(::Type{<:SpectralData}, ::AbstractStatistic) = nothing
 preferred_units(::Type{<:SpectralData}, ::ChiSquared) = u"counts / (s * keV)"
 preferred_units(::Type{<:SpectralData}, ::Cash) = u"counts"
 
+_as_unit(::Unitful.FreeUnits{U}) where {U} = first(U)
+
+function _adjust_by_unit_difference!(
+    ΔE,
+    exposure_time,
+    x,
+    ::Unitful.FreeUnits{Units},
+) where {Units}
+    foreach(Units) do u
+        if u === _as_unit(u"s")
+            @. x = x * exposure_time
+        elseif u === _as_unit(u"s^-1")
+            @. x = x / exposure_time
+        elseif u === _as_unit(u"keV")
+            @. x = x * ΔE
+        elseif u === _as_unit(u"keV^-1")
+            @. x = x / ΔE
+        else
+            throw(
+                ErrorException(
+                    "Cannot adjust units: unhandled unit difference: $(u) (if you are expecting to be able to use specific units, please open a bug report).",
+                ),
+            )
+        end
+    end
+    nothing
+end
+
+function adjust_to_units!(data::SpectralData, s::Spectrum, x, units)
+    ΔE = bin_widths(data)
+    exposure_time = s.exposure_time
+    _adjust_by_unit_difference!(ΔE, exposure_time, x, units / s.units)
+    x
+end
+
+function set_units!(s::SpectralData, units)
+    adjust_to_units!(s, s.spectrum, s.spectrum.data, units)
+    adjust_to_units!(s, s.spectrum, s.spectrum.errors, units)
+    s.spectrum.units = units
+    if has_background(s)
+        adjust_to_units!(s, s.background, s.background.data, units)
+        adjust_to_units!(s, s.background, s.background.errors, units)
+        s.background.units = units
+    end
+    s
+end
+
 # internal methods
 
 function rebin_if_different_domains!(output, data_domain, model_domain, input)
@@ -478,10 +509,11 @@ macro _forward_SpectralData_api(args)
             SpectralFitting.set_domain!(getfield(t, $(field)), args...)
         SpectralFitting.error_statistic(t::$(T)) =
             SpectralFitting.error_statistic(getfield(t, $(field)))
+        SpectralFitting.set_units!(t::$(T), args...) =
+            SpectralFitting.set_units!(getfield(t, $(field)), args...)
     end |> esc
 end
 
-
 # printing utilities
 
 function Base.show(io::IO, @nospecialize(data::SpectralData))
@@ -576,4 +608,5 @@ export SpectralData,
     drop_channels!,
     normalize!,
     subtract_background!,
-    set_domain!
+    set_domain!,
+    set_units!
diff --git a/src/datasets/spectrum.jl b/src/datasets/spectrum.jl
index 00f6ce03..cfaee641 100644
--- a/src/datasets/spectrum.jl
+++ b/src/datasets/spectrum.jl
@@ -31,12 +31,10 @@ end
 supports(::Type{<:Spectrum}) = (ContiguouslyBinned(),)
 
 function make_objective(::ContiguouslyBinned, dataset::Spectrum)
-    check_units_warning(dataset.units)
     dataset.data
 end
 
 function make_objective_variance(::ContiguouslyBinned, dataset::Spectrum)
-    check_units_warning(dataset.units)
     dataset.errors .^ 2
 end
 
diff --git a/src/fitting/config.jl b/src/fitting/config.jl
index 3dd82bcd..e3f341a0 100644
--- a/src/fitting/config.jl
+++ b/src/fitting/config.jl
@@ -72,8 +72,11 @@ function make_multi_config(prob::FittingProblem, stat::AbstractStatistic)
         all(model -> implementation(model) isa JuliaImplementation, prob.model.m) ?
         JuliaImplementation() : XSPECImplementation()
 
-    layout = common_support(prob.model.m..., prob.data.d...)
-
+    # TODO: preferred units needs to be the same for all datasets?
+    layout = with_units(
+        common_support(prob.model.m..., prob.data.d...),
+        preferred_units(first(prob.data.d), stat),
+    )
     variances = map(d -> make_objective_variance(layout, d), prob.data.d)
     # build index mappings for pulling out the data
     domains, domain_mapping = _build_domain_mapping(layout, prob.data)
diff --git a/src/fitting/statistics.jl b/src/fitting/statistics.jl
index 654fbf5b..992d0a79 100644
--- a/src/fitting/statistics.jl
+++ b/src/fitting/statistics.jl
@@ -24,28 +24,18 @@ end
 function measure(::Cash, S, ŷ)
     2 * sum(@.(ŷ - S + S * (log(S) - log(ŷ))))
 end
-function measure(s::Cash, config::FittingConfig, y, ŷ, σ²)
-    exp_time = config.prob.data[1].data.spectrum.exposure_time
-    prefactor = exp_time .* bin_widths(config.prob.data[1].data)
-    measure(s, y .* prefactor, ŷ, σ²)
+function measure(s::Cash, slice::FittingResultSlice, u = slice.u)
+    ŷ = _invoke_and_transform!(get_cache(slice), slice.domain, u)
+    measure(s, slice.objective, ŷ)
 end
 
 function _f_wrap_objective(stat::Cash, config::FittingConfig)
     f = _f_objective(config)
-    exp_time = config.prob.data[1].spectrum.exposure_time
-    prefactor = exp_time .* bin_widths(config.prob.data[1])
     function _objective(parameters, domain)
         ŷ = f(domain, parameters)
-        measure(stat, config.objective .* prefactor, ŷ .* prefactor)
+        measure(stat, config.objective, ŷ)
     end
 end
 
-function measure(s::Cash, slice::FittingResultSlice, u = slice.u)
-    data = get_dataset(slice)
-    exp_time = data.data.spectrum.exposure_time
-    prefactor = exp_time .* bin_widths(data.data)
-    ŷ = _invoke_and_transform!(get_cache(slice), slice.domain, u)
-    measure(s, slice.objective .* prefactor, y .* prefactor̂)
-end
 
 export ChiSquared, Cash, measure
diff --git a/src/simulate.jl b/src/simulate.jl
index 90fee63b..af0b18e6 100644
--- a/src/simulate.jl
+++ b/src/simulate.jl
@@ -109,7 +109,7 @@ function _make_simulation_fitting_config(
         model,
         input_domain,
         objective,
-        _fold_transformer(T, layout, R, ΔE, input_domain),
+        _fold_transformer(T, one(eltype(ΔE)), layout, R, ΔE, input_domain),
     )
 
     free_params = collect(filter(isfree, parameter_tuple(model)))