Number-weighted mdf_amr

src/fitting/mdf.jl

@@ -34,57 +34,47 @@ function mdf_amr(coeffs::AbstractVector{<:Number}, # Stellar mass coefficients
     return unique_MH[p], mass_mdf[p]
 end
 
-# function mdf_amr(stellar_masses::AbstractVector{<:Number},
-#                  logAge::AbstractVector{<:Number},
-#                  metallicities::AbstractVector{<:Number},
-#                  relweights::AbstractVector{<:Number};
-#                  relweightsmin::Number=0)
-
-#     unique_logAge = unique(logAge)
-#     unique_MH = unique(metallicities)
-#     @assert length(stellar_masses) == length(unique_logAge) # one SFR / stellar mass coefficient per unique logAge
-#     @assert length(logAge) == length(metallicities) == length(relweights)
-#     @assert all(x -> x ≥ 0, relweights) # All relative weights must be \ge 0
-#     @assert relweightsmin >= 0 # By definition relweightsmin must be greater than or equal to 0
-
-#     # Identify which of the provided models are significant enough
-#     # to be included in the fitting on the basis of the provided `relweightsmin`.
-#     if relweightsmin != 0
-#         keep_idx = truncate_relweights(relweightsmin, relweights, logAge) # Method defined below
-#         models = models[keep_idx]
-#         logAge = logAge[keep_idx]
-#         metallicities = metallicities[keep_idx]
-#         relweights = relweights[keep_idx]
-#     end
-
-#     # Loop through all unique logAge entries and ensure sum over relweights = 1
-#     for la in unique_logAge
-#         good = findall( logAge .== la )
-#         goodsum = sum(relweights[good])
-#         if !(goodsum ≈ 1)
-#             # Don't warn if relweightsmin != 0 as it will ALWAYS need to renormalize in this case
-#             if relweightsmin == 0
-#                 @warn "The relative weights for logAge="*string(la)*" provided to `fixed_amr` do not sum to 1 and will be renormalized in place. This warning is suppressed for additional values of logAge." maxlog=1
-#             end
-#             relweights[good] ./= goodsum
-#         end
-#     end
-
-#     # Make scratch array for holding full coefficient vector
-#     coeffs = similar(logAge)
+"""
+    (unique_MH, mass_mdf) =
+    mdf_amr(coeffs::AbstractVector{<:Number},
+            logAge::AbstractVector{<:Number},
+            metallicities::AbstractVector{<:Number},
+            models::Union{AbstractVector{<:AbstractMatrix{<:Number}},
+                          AbstractMatrix{<:Number}})
 
-#     # Compute the index masks for each unique entry in logAge so we can
-#     # construct the full coefficients vector when evaluating the likelihood
-#     idxlogAge = [findall( ==(i), logAge) for i in unique_logAge]
+Calculates the number-weighted metallicity distribution function given a set of coefficients `coeffs` for stellar populations with logarithmic ages `logAge=log10(age [yr])`, metallicities given by `metallicities`, and Hess diagram templates `models`. This function constructs a composite Hess diagram (see [`composite!`](@ref StarFormationHistories.composite!) for definition) out of the `coeffs` and `models` that match each unique entry in `metallicites`. The sums over the bins of these composite Hess diagrams then give the total predicted number of stars in the Hess diagram viewport for each metallicity. The raw number counts for each unique metallicity are returned as `mass_mdf` -- these can be renormalized afterward to sum to one by calculating `mass_mdf ./ sum(mass_mdf)`.
 
-#     # Expand the SFH coefficients into per-model coefficients
-#     for (i, idxs) in enumerate(idxlogAge)
-#         @inbounds coeffs[idxs] .= relweights[idxs] .* stellar_masses[i]
-#     end
+# Examples
+```jldoctest; setup = :(import StarFormationHistories: mdf_amr)
+julia> mdf_amr([1.0, 2.0, 1.0], [10, 10, 10], [-2, -1.5, -1],
+               map(Base.Fix2(fill, (100, 100)), [1.0, 2.0, 1.0]))
+([-2.0, -1.5, -1.0], [10000.0, 40000.0, 10000.0])
+```
+"""
+function mdf_amr(coeffs::AbstractVector{<:Number},
+                 logAge::AbstractVector{<:Number},
+                 metallicities::AbstractVector{<:Number},
+                 models::AbstractMatrix{T}) where T <: Number # For stacked models
 
-#     # Now, loop through and sum all the coeffs for each unique metallicity
-#     # This will form an (unnormalized) mass-weighted MDF.
-#     mass_mdf = [sum(coeffs[idx]) for idx in (findall( ==(i), metallicities) for i in unique_MH)]
-#     return unique_MH, mass_mdf
+    @assert length(coeffs) == length(logAge) == length(metallicities) == size(models, 2)
+    # Loop through all unique metallicities and sum the number of stars predicted
+    # given the stellar mass / SFR coefficients `coeffs` and the CMD models in `models`.
+    unique_MH = unique(metallicities)
+    mass_mdf = similar(unique_MH)
+    composite = Vector{T}(undef, size(models, 1))
+    for (i, mh) in enumerate(unique_MH)
+        idxs = findall(==(mh), metallicities)
+        # mul!(composite, view(models,:,idxs), view(coeffs, idxs))
+        mul!(composite, models[:,idxs], coeffs[idxs]) # More memory but faster mul! than views
+        mass_mdf[i] = sum(composite) # VectorizationBase.vsum not much faster
+    end
+    p = sortperm(unique_MH) # Return in sorted order of unique_MH
+    return unique_MH[p], mass_mdf[p]
+end
 
-# end
+# Just stack models and call above function
+mdf_amr(coeffs::AbstractVector{<:Number},
+        logAge::AbstractVector{<:Number},
+        metallicities::AbstractVector{<:Number},
+        models::AbstractVector{<:AbstractMatrix{<:Number}}) =
+            mdf_amr(coeffs, logAge, metallicities, stack_models(models))