Refactor AMR code to use new generic interface

docs/make.jl

@@ -42,9 +42,11 @@ makedocs(
                  ["fitting/fitting_intro.md",
                   "fitting/unconstrained.md",
                   "Constrained Metallicity Evolution" => 
-                      ["fitting/hierarchical/linear_amr.md",
-                       "fitting/hierarchical/log_amr.md",
-                       "fitting/hierarchical/fixed_amr.md",
+                      ["fitting/hierarchical/overview.md",
+                       "AMRs" =>
+                           ["fitting/hierarchical/linear_amr.md",
+                            "fitting/hierarchical/log_amr.md",
+                            "fitting/hierarchical/fixed_amr.md"],
                        "MZRs" =>
                            ["fitting/hierarchical/MZR/MZR.md"],
                        "fitting/hierarchical/dispersion_models.md"],

docs/src/fitting/hierarchical/MZR/MZR.md

@@ -22,11 +22,7 @@ We provide a generic interface for describing the analytic form of the MZR so th
 PowerLawMZR
 ```
 
-The per-SSP stellar mass coefficients (``r_{j,k}`` in the [derivation](@ref mzr_derivation)) can be derived from an MZR model, a [metallicity dispersion model](@ref dispersion_models), the per-unique-log(age) stellar mass coefficients (``R_j`` in the [derivation](@ref mzr_derivation)), and the set of SSP logarithmic ages `logAge = log10(age [yr])` and metallicites using [`calculate_coeffs`](@ref),
-
-```@docs
-calculate_coeffs(::StarFormationHistories.AbstractMZR, ::StarFormationHistories.AbstractDispersionModel, ::AbstractVector{<:Number}, ::AbstractVector{<:Number}, ::AbstractVector{<:Number})
-```
+The per-SSP stellar mass coefficients (``r_{j,k}`` in the [derivation](@ref mzr_derivation)) can be derived from an MZR model, a [metallicity dispersion model](@ref dispersion_models), the per-unique-log(age) stellar mass coefficients (``R_j`` in the [derivation](@ref mzr_derivation)), and the set of SSP logarithmic ages `logAge = log10(age [yr])` and metallicites using [`calculate_coeffs`](@ref StarFormationHistories.calculate_coeffs).
 
 ## [Mass-Metallicity Relation API](@id mzr_API)
 
@@ -42,6 +38,9 @@ StarFormationHistories.transforms(::StarFormationHistories.AbstractMZR)
 StarFormationHistories.free_params(::StarFormationHistories.AbstractMZR)
 ```
 
+## Fitting and Sampling Methods
+The MZRs listed above support the generic fitting and sampling methods listed in the [overview](@ref metal_evo_intro).
+
 ## [Derivation](@id mzr_derivation)
 
 We once again wish to derive the gradient of the objective function with respect to the fitting parameters to enable gradient-based optimization. Our derivation will reuse some of the notation developed in the section on the [linear AMR](@ref linear_amr_section). The main difference is that instead of expressing the mean metallicity as a function of time ``\langle [\text{M}/\text{H}] \rangle (t)``, with an MZR we express the mean metallicity as a function of stellar mass at that time ``\langle [\text{M}/\text{H}] \rangle (\text{M}_*(t))``. This means that the partial derivatives of the objective with respect to the stellar mass coefficients have a more complex form than in the AMR case, as changing the stellar mass formed 10 Gyr ago (for example) would change the total stellar mass at all more recent times, which in turn changes the mean metallicity expected at all more recent times.

docs/src/fitting/hierarchical/linear_amr.md

@@ -1,14 +1,6 @@
 # [Linear Age-Metallicity Relation](@id linear_amr_section)
 
-## [Why Should Metallicity Evolutions Be Constrained?](@id metal_evo_intro)
-
-While the above methods work well for optimizing the per-template ``r_j`` as a means for fitting SFHs, these methods can produce metallicity evolutions that could be considered unphysical, with large changes in the mean metallicity over small changes in time. An example of this type of behavior is shown in the SFH fit below.
-
-![Example of a SFH fit with variations in the metallicity evolution.](figures/mean_mh.png)
-
-While some metallicity variation in the star-forming gas is to be expected, these variations in the SFH fit can end up being quite large depending on the data and isochrone grid adopted. A solution is to construct a more physically-motivated model.
-
-We can do this using a hierarchical model with a parameterized metallicity distribution function (MDF) where the the ``r_j`` are not the parameters directly optimized. Rather, we can optimize one stellar mass (or star formation rate) parameter per age bin, and then a number of MDF parameters that determine how that stellar mass is split between models with different metallicities at fixed age. An example for one such MDF model is a linear mean metallicity relation ``\langle [\text{M}/\text{H}] \rangle (t) = \alpha \, \left( T_\text{max} - t \right) + \beta`` with a Gaussian distribution in metallicity at fixed age. ``T_\text{max}`` here is the earliest lookback time under consideration such that ``\langle [\text{M}/\text{H}] \rangle (T_\text{max}) = \beta``. If the per-age-bin stellar mass coefficients are ``R_j``, the age of the stellar population ``j`` is ``t_j``, and the metallicity of population ``k`` is ``[\text{M}/\text{H}]_k``, then we can write the per-model ``r_{j,k}`` (where we are now using separate indices for age and metallicity) as
+Here we describe the linear age-metallicity relation ``\langle [\text{M}/\text{H}] \rangle (t) = \alpha \, \left( T_\text{max} - t \right) + \beta`` with a Gaussian distribution in metallicity at fixed age as described by the [`GaussianDispersion`](@ref) dispersion model. ``T_\text{max}`` here is the earliest lookback time under consideration such that ``\langle [\text{M}/\text{H}] \rangle (T_\text{max}) = \beta``. If the per-age-bin stellar mass coefficients are ``R_j``, the age of the stellar population ``j`` is ``t_j``, and the metallicity of population ``k`` is ``[\text{M}/\text{H}]_k``, then we can write the per-model ``r_{j,k}`` (where we are now using separate indices for age and metallicity) as
 
 ```math
 \begin{aligned}
@@ -23,39 +15,28 @@ where the numerator is the MDF at fixed age evaluated at metallicity ``[\text{M}
 m_i = \sum_{j,k} \, r_{j,k} \; c_{i,j,k}
 ```
 
-Below we show a fit using this hierarchical model to the same data as above. 
+Below we show a fit using this hierarchical model to the same data as was used to derive the unconstrained fit in [the introduction](@ref metal_evo_intro). 
 
 ![Example of a SFH fit with a linear metallicity evolution.](figures/mdf_model.png)
 
-## Fitting Functions
-
-We provide the method [`StarFormationHistories.fit_templates_mdf`](@ref) to fit this model to an observed Hess diagram.
+This model is represented by the [`LinearAMR`](@ref) type, which is a subtype of [`AbstractAMR`](@ref StarFormationHistories.AbstractAMR).
 
 ```@docs
-StarFormationHistories.fit_templates_mdf
-StarFormationHistories.LogTransformMDFσResult
-StarFormationHistories.LogTransformMDFResult
+StarFormationHistories.AbstractAMR
+LinearAMR
 ```
 
-The method [`StarFormationHistories.construct_x0_mdf`](@ref) can be used to construct the stellar mass components ``R_j`` of the initial guess vector `x0`
+## Fitting Functions
 
-```@docs
-StarFormationHistories.construct_x0_mdf
-```
+[`fit_sfh`](@ref) and [`sample_sfh`](@ref) both work with this AMR model.
 
-and [`StarFormationHistories.calculate_coeffs_mdf`](@ref) can be used to calculate per-template stellar mass coefficients (the ``r_{j,k}`` above) given the results of a fit (which will be the ``R_j`` in the equations above)
+The method [`StarFormationHistories.construct_x0_mdf`](@ref) can be used to construct the stellar mass components ``R_j`` of the initial guess vector `x0`.
 
 ```@docs
-StarFormationHistories.calculate_coeffs_mdf
+StarFormationHistories.construct_x0_mdf
 ```
 
-## Sampling Methods
-
-We additionally offer a sampling method for this linear age-metallicity relation using HMC:
-
-```@docs
-StarFormationHistories.hmc_sample_mdf
-```
+and [`calculate_coeffs`](@ref) can be used to calculate per-template stellar mass coefficients (the ``r_{j,k}`` above) given the results of a fit (which will be the ``R_j`` in the equations above).
 
 ## [Implementation](@id linear_amr_implementation)
 

docs/src/fitting/hierarchical/log_amr.md

@@ -11,7 +11,15 @@ This model differs from the [linear age-metallicity relation (AMR)](@ref linear_
 \end{aligned}
 ```
 
-with ``T_\text{max}`` being the earliest lookback time under consideration, such that ``\langle Z \rangle (T_\text{max}) = \beta``. We choose this parameterization so that positive ``\alpha`` and ``\beta`` result in an age-metallicity relation that is monotonically increasing with decreasing lookback time ``t``. We model the spread in metallicities at fixed ``t`` as Gaussian in [M/H], identically to how it is modelled in the linear AMR case. This implies the spread is asymmetric in ``Z``; this can be seen in the output of `examples/log_amr/log_amr_example.jl`, shown below, which illustrates the relative weights due to a logarithmic AMR across a grid of ages and metallicities. The per-model coefficients implied by a such a logarithmic AMR can be calculated with [`calculate_coeffs_logamr`](@ref StarFormationHistories.calculate_coeffs_logamr).
+with ``T_\text{max}`` being the earliest lookback time under consideration, such that ``\langle Z \rangle (T_\text{max}) = \beta``. We choose this parameterization so that positive ``\alpha`` and ``\beta`` result in an age-metallicity relation that is monotonically increasing with decreasing lookback time ``t``. If we model the spread in metallicities at fixed ``t`` as Gaussian in [M/H] with the [`GaussianDispersion`](@ref) dispersion model, this implies the spread is asymmetric in ``Z``; this can be seen in the output of `examples/log_amr/log_amr_example.jl`, shown below, which illustrates the relative weights due to a logarithmic AMR across a grid of ages and metallicities.
+
+This logarithmic AMR is implemented with the [`LogarithmicAMR`](@ref) type, which is a subtype of [`AbstractAMR`](@ref StarFormationHistories.AbstractAMR).
+
+```@docs
+LogarithmicAMR
+```
+
+The per-model coefficients (the ``r_{j,k}`` above) implied by a such a logarithmic AMR can be calculated with [`calculate_coeffs`](@ref) and an example is shown below.
 
 ```@example
 ENV["GKSwstype"] = "100" # https://discourse.julialang.org/t/generation-of-documentation-fails-qt-qpa-xcb-could-not-connect-to-display/60988 # hide
@@ -21,37 +29,11 @@ savefig("log_amr_plot.svg"); nothing # hide
 
 ![Visualization of the relative weights across a grid of logAge and metallicity under a logarithmic age-metallicity relation.](log_amr_plot.svg)
 
-```@docs
-StarFormationHistories.calculate_coeffs_logamr
-```
-
 ## Fitting Functions
 
-The main function we provide to fit star formation histories to Hess diagrams under the logarithmic age-metallicity relation is [`fit_templates_logamr`](@ref StarFormationHistories.fit_templates_logamr). This function operates similarly to the fitting function for the linear AMR model, [`fit_templates_mdf`](@ref StarFormationHistories.fit_templates_mdf). 
-
-```@docs
-StarFormationHistories.fit_templates_logamr
-```
-
-## Sampling Functions
+[`fit_sfh`](@ref) and [`sample_sfh`](@ref) both work with this AMR model.
 
-```@docs
-StarFormationHistories.hmc_sample_logamr
-```
-
-## Fixed Logarithmic Age-Metallicity Relation
-
-We support fitting only the star formation parameters by adopting fixed values for ``\alpha``, ``\beta``, and ``\sigma`` through the [`fixed_log_amr`](@ref StarFormationHistories.fixed_log_amr) method.
-
-```@docs
-StarFormationHistories.fixed_log_amr
-```
-
-We provide the [calculate\_αβ\_logamr](@ref StarFormationHistories.calculate_αβ_logamr) convenience function to calculate the slope ``\alpha`` and intercept ``\beta`` from two points on the age-metallicity relation.
-
-```@docs
-StarFormationHistories.calculate_αβ_logamr
-```
+The method [`StarFormationHistories.construct_x0_mdf`](@ref) can be used to construct the stellar mass components ``R_j`` of the initial guess vector `x0`.
 
 ## Implementation
 

docs/src/fitting/hierarchical/overview.md

@@ -0,0 +1,49 @@
+# [Overview](@id metal_evo_intro)
+
+Why should the metallicity evolution be constrained? While the above methods work well for optimizing the per-template ``r_j`` as a means for fitting SFHs, these methods can produce metallicity evolutions that could be considered unphysical, with large changes in the mean metallicity over small changes in time. An example of this type of behavior is shown in the SFH fit below.
+
+![Example of a SFH fit with variations in the metallicity evolution.](figures/mean_mh.png)
+
+While some metallicity variation in the star-forming gas is to be expected, these variations in the SFH fit can end up being quite large depending on the data and isochrone grid adopted. A solution is to construct a more physically-motivated model.
+
+We can do this using a hierarchical model with a parameterized metallicity evolution where the the ``r_j`` are not the parameters directly optimized. Rather, we can optimize one stellar mass (or star formation rate) parameter per age bin, and then a number of metallicity evolution parameters that determine how that stellar mass is split between models with different metallicities at fixed age.
+
+In most star formation history analyses, the metallicities are constrained through *age-metallicity relations (AMRs)*, where the mean metallicity at time ``t`` is a function of time and a small set of metallicity evolution parameters. A popular AMR model is the linear age-metallicity relation ``\langle [\text{M}/\text{H}] \rangle (t) = \alpha \, \left( T_\text{max} - t \right) + \beta`` with a Gaussian distribution in metallicity at fixed age. ``T_\text{max}`` here is the earliest lookback time under consideration such that ``\langle [\text{M}/\text{H}] \rangle (T_\text{max}) = \beta``. This model is described in more detail [here](@ref linear_amr_section).
+
+AMRs have historically been popular because they are generally capable of producing reasonable fits to observed data and it is relatively easy to derive the gradient of the objective function with respect to the AMR parameters analytically. However, in AMR models there is no direct link between the SFRs being fit and the metallicity evolution as a function of time, even though the two should in principle have some correlation as stellar processes are responsible for enriching the ISM.
+
+A promising avenue of research involves fitting *mass-metallicity relations* (MZRs) rather than AMRs. In these models, the mean metallicity of stars forming at time ``t`` is a function of the total stellar mass of the population at that time -- therefore, the mean metallicity evolution changes self-consistently with the SFRs during the fitting process, resulting in a metallicity evolution that is meaningfully coupled to the star formation history. Additionally, AMRs can be difficult to compare between different galaxies because they do not reflect the different SFHs of the galaxies, whereas MZRs can be compared between galaxies much more easily. Our methods for MZR fitting are described in more detail [here](@ref MZR).
+
+## Generic Methods
+
+While there are some methods in this package that are unique to AMR or MZR models, we present a minimal unified interface that can be used to fit SFHs under both types of models. To support multiple dispatch, we define [`AbstractMetallicityModel`](@ref StarFormationHistories.AbstractMetallicityModel) as the abstract supertype of [`AbstractAMR`](@ref StarFormationHistories.AbstractAMR) and [`AbstractMZR`](@ref StarFormationHistories.AbstractMZR), which are each the supertypes for AMR and MZR types, respectively.
+
+```@docs
+StarFormationHistories.AbstractMetallicityModel
+```
+
+The generic methods that can be used for both AMRs and MZRs are described here. The main method for obtaining best-fit star formation histories is [`fit_sfh`](@ref).
+
+```@docs
+fit_sfh
+```
+
+This function returns an instance of [`CompositeBFGSResult`](@ref StarFormationHistories.CompositeBFGSResult).
+
+```@docs
+StarFormationHistories.CompositeBFGSResult
+StarFormationHistories.BFGSResult
+```
+
+This can be used to obtain random samples under a multivariable Normal approximation to the posterior or used to initialize a Hamiltonian Monte Carlo (HMC) sampling process to obtain more accurate posterior samples with [`sample_sfh`](@ref) and its multi-threaded alternative [`tsample_sfh`](@ref StarFormationHistories.tsample_sfh).
+
+```@docs
+sample_sfh
+StarFormationHistories.tsample_sfh
+```
+
+The per-SSP stellar mass coefficients (``r_{j,k}`` in the [derivation](@ref mzr_derivation)) can be derived from a metallicity model, a [metallicity dispersion model](@ref dispersion_models), the per-unique-log(age) stellar mass coefficients (``R_j`` in the [derivation](@ref mzr_derivation)), and the set of SSP logarithmic ages `logAge = log10(age [yr])` and metallicites using [`calculate_coeffs`](@ref StarFormationHistories.calculate_coeffs). Alternatively a [`CompositeBFGSResult`](@ref StarFormationHistories.CompositeBFGSResult) can be fed into this method and the first three arguments will be read from the result object.
+
+```@docs
+calculate_coeffs
+```

docs/src/helpers.md

@@ -28,6 +28,7 @@ StarFormationHistories.X_from_Z
 StarFormationHistories.MH_from_Z
 StarFormationHistories.dMH_dZ
 StarFormationHistories.Z_from_MH
+StarFormationHistories.dZ_dMH
 StarFormationHistories.mdf_amr
 ```
 

src/StarFormationHistories.jl

@@ -9,7 +9,7 @@ using IrrationalConstants: logten
 import LBFGSB # Used for one method in fitting.jl
 import LineSearches # For configuration of Optim.jl
 # Need mul! for composite!, ∇loglikelihood!;
-using LinearAlgebra: diag, Hermitian, hermitianpart, mul!
+using LinearAlgebra: BLAS, diag, Hermitian, hermitianpart, mul!
 import LogDensityProblems # For interfacing with DynamicHMC
 using LoopVectorization: @turbo
 import LoopVectorization: can_turbo # Extending for our functions
@@ -904,6 +904,8 @@ export generate_stars_mass, generate_stars_mag, generate_stars_mass_composite,
 export fit_templates,
     hmc_sample,
     GaussianDispersion,
+    LinearAMR,
+    LogarithmicAMR,
     PowerLawMZR,
     fit_sfh,
     sample_sfh,