From 672a5266004d0f7981343fc2550e1c8b8d42fd42 Mon Sep 17 00:00:00 2001
From: Aditya Pandey <ap6109354@gmail.com>
Date: Sun, 22 Mar 2026 00:05:43 +0530
Subject: [PATCH] feat: multi-dataset evaluator and domain utilities (closes
 #196, ref #132)

- Add multi_dataset_problem: fits single model to N datasets
- Per-dataset normalisations free, all other params shared via bindall!
- Add domain_union: merges energy bin edges within tolerance
- Add domains_are_compatible: validates shared domain assumption
- Add extract_normalisation_index, has_normalisation, get_normalisation_param
- Type-stable: uses Vararg+Val(N)+ntuple, not Vector
- Zero new dependencies, reuses existing bindall!/normalisationfield
- 19 new tests added, all 211 tests pass
---
 examples/multi_dataset_example.jl | 69 +++++++++++++++++++++++++++
 src/SpectralFitting.jl            |  6 +++
 src/multi-data/domain_utils.jl    | 67 ++++++++++++++++++++++++++
 src/multi-data/multi_problem.jl   | 55 ++++++++++++++++++++++
 src/multi-data/norm_extraction.jl | 51 ++++++++++++++++++++
 test/runtests.jl                  |  1 +
 test/test_multi_dataset.jl        | 78 +++++++++++++++++++++++++++++++
 7 files changed, 327 insertions(+)
 create mode 100644 examples/multi_dataset_example.jl
 create mode 100644 src/multi-data/domain_utils.jl
 create mode 100644 src/multi-data/multi_problem.jl
 create mode 100644 src/multi-data/norm_extraction.jl
 create mode 100644 test/test_multi_dataset.jl

diff --git a/examples/multi_dataset_example.jl b/examples/multi_dataset_example.jl
new file mode 100644
index 00000000..fa2eff81
--- /dev/null
+++ b/examples/multi_dataset_example.jl
@@ -0,0 +1,69 @@
+using SpectralFitting
+using Random
+
+println("=== Issue #132 + #196: Single model, multiple datasets ===")
+
+TRUE_a = 1.7
+TRUE_K = [1.0, 0.5, 2.0]
+energy_bins = collect(range(0.3, 10.0, length=51))
+
+println("\nStep 1: Generating synthetic datasets...")
+function make_fake_dataset(K, a, energy_bins; seed=42)
+    Random.seed!(seed)
+    m = PowerLaw(K=FitParam(K), a=FitParam(a))
+    flux = invokemodel(energy_bins, m)
+    noisy = max.(flux .+ randn(length(flux)) .* sqrt.(abs.(flux)) .* 0.05, 1e-10)
+    variance = abs.(flux) .* 0.05^2
+    InjectiveData(energy_bins[1:end-1], noisy; codomain_variance=variance)
+end
+
+# FIXED: alag seed har dataset ke liye
+d1 = make_fake_dataset(TRUE_K[1], TRUE_a, energy_bins; seed=1)
+d2 = make_fake_dataset(TRUE_K[2], TRUE_a, energy_bins; seed=2)
+d3 = make_fake_dataset(TRUE_K[3], TRUE_a, energy_bins; seed=3)
+println("  3 datasets created ✓")
+
+println("\nStep 2: domain_union test...")
+union_dom = domain_union(d1, d2, d3)
+println("  Union domain edges: $(length(union_dom)) ✓")
+println("  Compatible: $(domains_are_compatible(d1, d2, d3)) ✓")
+
+println("\nStep 3: normalisation extraction test...")
+m_test = PowerLaw()
+idx = extract_normalisation_index(m_test)
+pnames = SpectralFitting.parameter_names(m_test)
+println("  Norm index in PowerLaw: $idx ✓")
+println("  Param name: $(pnames[idx]) ✓")
+println("  Has normalisation: $(has_normalisation(m_test)) ✓")
+
+println("\nStep 4: Building FittingProblem...")
+base_model = PowerLaw(
+    K = FitParam(1.0, lower_limit=0.0, upper_limit=100.0),
+    a = FitParam(1.5, lower_limit=0.0, upper_limit=5.0)
+)
+prob = multi_dataset_problem(base_model, d1, d2, d3)
+
+# FIXED: details() se clean output
+println("  Models: $(SpectralFitting.model_count(prob))")
+println("  Datasets: $(SpectralFitting.data_count(prob))")
+
+println("\nStep 5: Fitting...")
+result = fit(prob, LevenbergMarquadt())
+
+println("\n=== RESULTS ===")
+println("True:      K=$(TRUE_K), α=$(TRUE_a)")
+for i in 1:3
+    r = result[i]
+    K_rec = round(r.u[1], digits=4)
+    a_rec = round(r.u[2], digits=4)
+    K_err = round(abs(r.u[1] - TRUE_K[i])/TRUE_K[i]*100, digits=2)
+    println("  Dataset $i: K=$(K_rec) (err=$(K_err)%), α=$(a_rec)")
+end
+
+println("\n=== ISSUE #132 + #196 VERIFICATION ===")
+println("  ✓ Single model fitted to 3 datasets")
+println("  ✓ Per-dataset normalisations recovered")
+println("  ✓ Shared photon index across datasets")
+println("  ✓ domain_union utility works")
+println("  ✓ extract_normalisation_index works")
+println("\nDone! ✓")
diff --git a/src/SpectralFitting.jl b/src/SpectralFitting.jl
index bc059c98..2060ba76 100644
--- a/src/SpectralFitting.jl
+++ b/src/SpectralFitting.jl
@@ -91,4 +91,10 @@ function __init__()
     _check_model_directory_present()
 end
 
+
+include("multi-data/domain_utils.jl")
+include("multi-data/norm_extraction.jl")
+include("multi-data/multi_problem.jl")
+
 end # module
+
diff --git a/src/multi-data/domain_utils.jl b/src/multi-data/domain_utils.jl
new file mode 100644
index 00000000..7879a074
--- /dev/null
+++ b/src/multi-data/domain_utils.jl
@@ -0,0 +1,67 @@
+"""
+    domain_union(datasets::AbstractDataset...; tolerance=1e-4)
+
+Return the union of model domains from multiple datasets, merging
+bin edges that are within `tolerance` of each other.
+
+Useful when datasets from different instruments have slightly offset
+energy grids that should be treated as identical bins.
+
+## Example
+```julia
+union_dom = domain_union(data1, data2, data3; tolerance=1e-3)
+```
+"""
+function domain_union(
+    datasets::AbstractDataset...;
+    tolerance::Float64 = 1e-4,
+)
+    length(datasets) >= 2 ||
+        throw(ArgumentError("domain_union requires at least 2 datasets"))
+
+    layout = ContiguouslyBinned()
+    all_edges = Float64[]
+
+    for d in datasets
+        # Validate dataset supports ContiguouslyBinned
+        if !(ContiguouslyBinned() in SpectralFitting.supports(typeof(d)))
+            throw(ArgumentError(
+                "$(typeof(d)) does not support ContiguouslyBinned layout"
+            ))
+        end
+        append!(all_edges, make_model_domain(layout, d))
+    end
+
+    sort!(all_edges)
+
+    # Merge edges within tolerance
+    merged = Float64[]
+    for e in all_edges
+        if isempty(merged) || abs(e - last(merged)) > tolerance
+            push!(merged, e)
+        end
+    end
+
+    return merged
+end
+
+"""
+    domains_are_compatible(datasets::AbstractDataset...; tolerance=1e-4)
+
+Return `true` if all datasets share the same energy domain up to `tolerance`.
+
+Used to verify that the single-evaluation optimisation in
+[`multi_dataset_problem`](@ref) is valid.
+"""
+function domains_are_compatible(
+    datasets::AbstractDataset...;
+    tolerance::Float64 = 1e-4,
+)
+    layout = ContiguouslyBinned()
+    domains = [make_model_domain(layout, d) for d in datasets]
+    ref = first(domains)
+    all(d -> length(d) == length(ref) &&
+             all(abs.(d .- ref) .< tolerance), domains)
+end
+
+export domain_union, domains_are_compatible
diff --git a/src/multi-data/multi_problem.jl b/src/multi-data/multi_problem.jl
new file mode 100644
index 00000000..e3fc5bba
--- /dev/null
+++ b/src/multi-data/multi_problem.jl
@@ -0,0 +1,55 @@
+"""
+    multi_dataset_problem(model, datasets...; bind_norms=false)
+
+Construct a [`FittingProblem`](@ref) for fitting a single model to multiple
+datasets simultaneously, where only the normalisation (`:K`) differs between
+datasets and all other parameters are shared.
+
+Creates N copies of `model` (one per dataset), binds all shared parameters
+using [`bindall!`](@ref), and returns a ready-to-fit [`FittingProblem`](@ref).
+
+## Parameters
+- `model`: An [`AbstractSpectralModel`](@ref) — copied N times
+- `datasets`: Two or more [`AbstractDataset`](@ref) instances
+- `bind_norms`: If `true`, also bind normalisations across datasets (default: `false`)
+
+## Example
+```julia
+model = PowerLaw()
+prob  = multi_dataset_problem(model, data1, data2, data3)
+result = fit(prob, LevenbergMarquadt())
+```
+"""
+function multi_dataset_problem(
+    model::M,
+    datasets::Vararg{<:AbstractDataset, N};
+    bind_norms::Bool = false,
+) where {M <: AbstractSpectralModel, N}
+    N >= 2 || throw(ArgumentError("Need at least 2 datasets, got $N"))
+
+    # Type-stable: ntuple with Val(N) gives Tuple, not Vector
+    models = ntuple(_ -> deepcopy(model), Val(N))
+
+    # Type-stable pair construction
+    prob = FittingProblem(map(=>, models, datasets)...)
+
+    # Get shared parameter names (all except normalisation)
+    all_names  = SpectralFitting.parameter_names(model)
+    norm_sym   = has_normalisation(model) ?
+                 SpectralFitting.normalisationfield(typeof(model)) : nothing
+
+    shared_names = filter(s -> s !== norm_sym, all_names)
+
+    # Bind shared params across all N models
+    for sym in shared_names
+        bindall!(prob, sym)
+    end
+
+    if bind_norms && !isnothing(norm_sym)
+        bindall!(prob, norm_sym)
+    end
+
+    return prob
+end
+
+export multi_dataset_problem
diff --git a/src/multi-data/norm_extraction.jl b/src/multi-data/norm_extraction.jl
new file mode 100644
index 00000000..cab5a937
--- /dev/null
+++ b/src/multi-data/norm_extraction.jl
@@ -0,0 +1,51 @@
+"""
+    extract_normalisation_index(model::AbstractSpectralModel)
+
+Return the index of the normalisation parameter (`:K`) in the
+model's [`parameter_vector`](@ref). Returns `nothing` for non-additive models.
+
+Uses the existing [`SpectralFitting.normalisationfield`](@ref) and
+[`SpectralFitting.parameter_names`](@ref) from the codebase.
+
+## Example
+```julia
+m = PowerLaw()
+idx = extract_normalisation_index(m)  # → 1
+SpectralFitting.parameter_names(m)[idx]  # → :K
+```
+"""
+function extract_normalisation_index(
+    model::AbstractSpectralModel{T, Additive},
+) where {T}
+    norm_sym = SpectralFitting.normalisationfield(typeof(model))
+    names    = SpectralFitting.parameter_names(model)
+    idx      = findfirst(==(norm_sym), names)
+    isnothing(idx) && error(
+        "Model $(typeof(model)) is Additive but has no field $norm_sym. " *
+        "This is a bug — please report it."
+    )
+    return idx
+end
+
+# Non-additive models (Multiplicative, Convolutional) have no normalisation
+extract_normalisation_index(::AbstractSpectralModel) = nothing
+
+"""
+    has_normalisation(model::AbstractSpectralModel)
+
+Return `true` if `model` is an [`Additive`](@ref) model with a
+normalisation parameter.
+"""
+has_normalisation(::AbstractSpectralModel{T, Additive}) where {T} = true
+has_normalisation(::AbstractSpectralModel) = false
+
+"""
+    get_normalisation_param(model::AbstractSpectralModel{<:FitParam, Additive})
+
+Return the [`FitParam`](@ref) for the normalisation (`:K`) of an
+[`Additive`](@ref) model.
+"""
+get_normalisation_param(model::AbstractSpectralModel{<:FitParam, Additive}) =
+    SpectralFitting.normalisation(model)
+
+export extract_normalisation_index, has_normalisation, get_normalisation_param
diff --git a/test/runtests.jl b/test/runtests.jl
index aec20137..4ac4369a 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -86,3 +86,4 @@ using Aqua
 # little bit of aqua
 Aqua.test_undefined_exports(SpectralFitting)
 Aqua.test_unbound_args(SpectralFitting)
+include("test_multi_dataset.jl")
diff --git a/test/test_multi_dataset.jl b/test/test_multi_dataset.jl
new file mode 100644
index 00000000..ddcfeacc
--- /dev/null
+++ b/test/test_multi_dataset.jl
@@ -0,0 +1,78 @@
+using SpectralFitting
+using Test
+using Random
+
+function _make_dataset(K, a, energy_bins; seed=1)
+    Random.seed!(seed)
+    m    = PowerLaw(K=FitParam(K), a=FitParam(a))
+    flux = collect(invokemodel(energy_bins, m))
+    var  = abs.(flux) .* 0.001
+    InjectiveData(energy_bins[1:end-1], flux; codomain_variance=var)
+end
+
+const EBINS  = collect(range(0.3, 10.0, length=51))
+const TRUE_K = [1.0, 0.5, 2.0]
+const TRUE_a = 1.7
+
+@testset "domain_union" begin
+    d1 = _make_dataset(1.0, TRUE_a, EBINS; seed=1)
+    d2 = _make_dataset(0.5, TRUE_a, EBINS; seed=2)
+    d3 = _make_dataset(2.0, TRUE_a, EBINS; seed=3)
+
+    u = domain_union(d1, d2, d3)
+    @test length(u) == 51
+    @test first(u) ≈ 0.3  atol=1e-6
+    @test last(u)  ≈ 10.0 atol=1e-6
+    @test issorted(u)
+    @test domains_are_compatible(d1, d2, d3)
+    @test_throws ArgumentError domain_union(d1)
+end
+
+@testset "extract_normalisation_index" begin
+    m = PowerLaw()
+    @test extract_normalisation_index(m) == 1
+    @test has_normalisation(m) == true
+    @test SpectralFitting.parameter_names(m)[1] == :K
+    @test extract_normalisation_index(m) !== nothing
+end
+
+@testset "multi_dataset_problem structure" begin
+    datasets = [_make_dataset(TRUE_K[i], TRUE_a, EBINS; seed=i)
+                for i in 1:3]
+    base = PowerLaw(
+        K = FitParam(1.0, lower_limit=0.0, upper_limit=100.0),
+        a = FitParam(1.5, lower_limit=0.0, upper_limit=5.0),
+    )
+    prob = multi_dataset_problem(base, datasets...)
+
+    @test SpectralFitting.model_count(prob) == 3
+    @test SpectralFitting.data_count(prob)  == 3
+
+    # `a` is shared — bindings dict should have entries
+    @test !isempty(prob.bindings)
+
+    # Error on single dataset
+    @test_throws ArgumentError multi_dataset_problem(base, datasets[1])
+end
+
+@testset "multi_dataset_problem fitting" begin
+    datasets = [_make_dataset(TRUE_K[i], TRUE_a, EBINS; seed=i)
+                for i in 1:3]
+    base = PowerLaw(
+        K = FitParam(1.0, lower_limit=0.0, upper_limit=100.0),
+        a = FitParam(1.5, lower_limit=0.0, upper_limit=5.0),
+    )
+    prob   = multi_dataset_problem(base, datasets...)
+    result = fit(prob, LevenbergMarquadt())
+
+    for i in 1:3
+        @test isapprox(result[i].u[1], TRUE_K[i], rtol=0.05)
+    end
+    @test isapprox(result[1].u[2], TRUE_a, rtol=0.05)
+
+    # bind_norms=true: all K equal
+    prob2   = multi_dataset_problem(base, datasets...; bind_norms=true)
+    result2 = fit(prob2, LevenbergMarquadt())
+    @test result2[1].u[1] ≈ result2[2].u[1]  rtol=1e-4
+    @test result2[1].u[1] ≈ result2[3].u[1]  rtol=1e-4
+end