diff --git a/Project.toml b/Project.toml index a0ea707..725dff2 100644 --- a/Project.toml +++ b/Project.toml @@ -14,10 +14,20 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" StatsModels = "3eaba693-59b7-5ba5-a881-562e759f1c8d" +[extras] +JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6" + +[weakdeps] +JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6" + +[extensions] +LassoJSONExt = "JSON" + [compat] DSP = "0.7, 0.8" Distributions = "0.25" GLM = "1.9" +JSON = "1" MLBase = "0.9" Reexport = "1" StatsBase = "0.33, 0.34" diff --git a/README.md b/README.md index 54cc6e8..008fc9c 100644 --- a/README.md +++ b/README.md @@ -76,6 +76,23 @@ tend to result in sparser coefficient estimates. More documentation is available at [![][docs-stable-img]][docs-stable-url]. +## Saving and loading a fitted model + +Loading `JSON.jl` alongside Lasso.jl enables +`write_json`/`read_json`, which save just the coefficients, intercept, and +link of a fitted `LassoModel`/`GammaLassoModel` — enough to reproduce +`predict` without keeping the full fit (and its underlying data) around: + +```julia +using Lasso, JSON + +m = fit(LassoModel, X, y, dist, link; select=MinAICc()) +write_json("model.json", m) + +m2 = read_json("model.json") +predict(m2, newX) ≈ predict(m, newX) +``` + ## TODO - User-specified weights are untested diff --git a/docs/make.jl b/docs/make.jl index 8a23af2..11f4c5a 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -17,6 +17,7 @@ makedocs( "Home" => "index.md", "Lasso paths" => "lasso.md", "Fused Lasso and trend filtering" => "smoothing.md", + "Saving and loading a fitted model" => "serialization.md", "Index" => "api.md", ], ) diff --git a/docs/src/serialization.md b/docs/src/serialization.md new file mode 100644 index 0000000..708085d --- /dev/null +++ b/docs/src/serialization.md @@ -0,0 +1,55 @@ +# Saving and loading a fitted model + +A fitted `RegularizedModel` (e.g. a `LassoModel`) can be reduced to the +coefficients, intercept flag, link, and offset requirement needed to call +`predict`, without keeping the full fit and its underlying data around. This +is done via [`InferenceModel`](@ref). + +`InferenceModel(m)` extracts this state from a fitted model (also works on +the `TableRegressionModel` wrapper returned by formula-based `fit`, as used +below). `to_dict` converts an `InferenceModel` to a plain `Dict`, and +`InferenceModel(d)` reconstructs one from such a `Dict`: + +```jldoctest +julia> using DataFrames, Lasso + +julia> data = DataFrame(X=[1,2,3], Y=[2,4,7]); + +julia> m = fit(LassoModel, @formula(Y ~ X), data; select=MinAICc()); + +julia> m2 = Lasso.InferenceModel(m); + +julia> d = Lasso.to_dict(m2); + +julia> m3 = Lasso.InferenceModel(d); + +julia> predict(m3, reshape(data.X, :, 1)) ≈ predict(m, data) +true +``` + +If the original model was fit with an `offset`, `predict(m2, newX; offset=...)` +requires it, matching `predict(m::RegularizedModel, newX; offset=...)`. +`InferenceModel` only supports a single fitted `LassoModel`/ +`GammaLassoModel` segment — calling it on a full `LassoPath`/`GammaLassoPath` +or on a `FusedLasso`/`TrendFilter` model raises an `ArgumentError`. + +This roundtrip doesn't depend on any particular serialization format. Loading +`JSON.jl` alongside Lasso.jl also enables `write_json`/`read_json`, which +build on `InferenceModel`/`to_dict` to save to and load from JSON directly +(and, like `InferenceModel`, accept a fitted model directly, not just an +`InferenceModel`): + +```julia +using Lasso, JSON + +write_json("model.json", m) +m2 = read_json("model.json") +predict(m2, newX) ≈ predict(m, newX) +``` + +```@docs +InferenceModel +to_dict +write_json +read_json +``` diff --git a/ext/LassoJSONExt.jl b/ext/LassoJSONExt.jl new file mode 100644 index 0000000..b6d1608 --- /dev/null +++ b/ext/LassoJSONExt.jl @@ -0,0 +1,27 @@ +module LassoJSONExt + +using Lasso +using Lasso: InferenceModel, to_dict +using JSON + +function Lasso.write_json(io::IO, m) + return JSON.print(io, to_dict(InferenceModel(m))) +end + +function Lasso.write_json(path::AbstractString, m) + open(path, "w") do io + Lasso.write_json(io, m) + end +end + +function Lasso.read_json(io::IO) + return InferenceModel(JSON.parse(io)) +end + +function Lasso.read_json(path::AbstractString) + open(path, "r") do io + Lasso.read_json(io) + end +end + +end diff --git a/src/Lasso.jl b/src/Lasso.jl index 2297431..0f74e2a 100644 --- a/src/Lasso.jl +++ b/src/Lasso.jl @@ -21,7 +21,8 @@ export RegularizationPath, LassoPath, GammaLassoPath, NaiveCoordinateDescent, minAICc, hasintercept, dof, aicc, distfun, linkfun, cross_validate_path, SegSelect, segselect, selectmodel, AllSeg, MinAIC, MinAICc, MinBIC, CVSegSelect, MinCVmse, MinCV1se, - RegularizedModel, LassoModel, GammaLassoModel + RegularizedModel, LassoModel, GammaLassoModel, + InferenceModel, to_dict, write_json, read_json ## HELPERS FOR SPARSE COEFFICIENTS @@ -693,6 +694,7 @@ end include("coordinate_descent.jl") include("gammalasso.jl") include("segselect.jl") +include("inference_model.jl") include("cross_validation.jl") include("deprecated.jl") diff --git a/src/inference_model.jl b/src/inference_model.jl new file mode 100644 index 0000000..f911bf2 --- /dev/null +++ b/src/inference_model.jl @@ -0,0 +1,141 @@ +""" +InferenceModel stores the minimal state of a fitted `RegularizedModel` needed to +call `predict`: coefficients, whether an intercept was fit, the link function, +and whether the model requires an `offset` at predict time. + +Constructed from a fitted model via `InferenceModel(m)` (works on a +`LassoModel`/`GammaLassoModel`, or its formula-fit `TableRegressionModel` +wrapper) or from a `Dict` via `InferenceModel(d)` (see [`to_dict`](@ref), or +[`read_json`](@ref)). +""" +struct InferenceModel + coef::Vector{Float64} + intercept::Bool + link::Link + hasoffset::Bool +end + +InferenceModel(m::InferenceModel) = m + +""" + InferenceModel(m::RegularizedModel) -> InferenceModel + InferenceModel(m::StatsModels.TableRegressionModel{<:RegularizedModel}) -> InferenceModel + +Extract the coefficients, intercept flag, link, and offset requirement of a +fitted `RegularizedModel` (e.g. a `LassoModel`), sufficient to call `predict` +without keeping the full fit and its underlying data around. +""" +function InferenceModel(m::RegularizedModel) + link = _get_link(m.lpm) + haskey(LINK_TYPE_TO_NAME, typeof(link)) || + error("Unsupported link type for serialization: $(typeof(link))") + return InferenceModel(coef(m), GLM.hasintercept(m.lpm), link, !isempty(m.lpm.rr.offset)) +end + +InferenceModel(mm::StatsModels.TableRegressionModel{<:RegularizedModel}) = InferenceModel(mm.model) + +InferenceModel(path::RegularizationPath) = throw(ArgumentError( + "InferenceModel is not defined for a full $(nameof(typeof(path))); it only supports a " * + "single selected segment. Use `fit(LassoModel, ...)` or `fit(GammaLassoModel, ...)` to " * + "fit and select one segment, then call `InferenceModel` on that.")) + +InferenceModel(m) = throw(ArgumentError( + "InferenceModel is not defined for $(typeof(m)); only a fitted LassoModel/GammaLassoModel " * + "segment (or its formula-fit TableRegressionModel wrapper) is supported.")) + +"link of the underlying GLM of a fitted RegularizedModel segment" +_get_link(lpm::LinearModel) = IdentityLink() +_get_link(lpm::GeneralizedLinearModel) = Link(lpm) + +# explicit type <-> name tables +# This deliberately does not use metaprogramming or `nameof` or similar, +# because if the julia type `IdentityLink` is renamed to e.g. `IdentityLink2`, +# we want to continue to read/write it as the string +# "IdentityLink" for stability in the serialized format. +const LINK_TYPE_TO_NAME = Dict( + IdentityLink => "IdentityLink", + LogitLink => "LogitLink", + LogLink => "LogLink", + InverseLink => "InverseLink", + CloglogLink => "CloglogLink", + CauchitLink => "CauchitLink", + ProbitLink => "ProbitLink", + SqrtLink => "SqrtLink", + InverseSquareLink => "InverseSquareLink", + NegativeBinomialLink => "NegativeBinomialLink", +) +const LINK_NAME_TO_CONSTRUCTOR = Dict(name => T for (T, name) in LINK_TYPE_TO_NAME) + +""" + to_dict(m::InferenceModel) -> Dict + +Convert an [`InferenceModel`](@ref) to a plain `Dict`, sufficient to +reconstruct it via `InferenceModel(d)`. +""" +function to_dict(m::InferenceModel) + linkname = get(LINK_TYPE_TO_NAME, typeof(m.link)) do + error("Unsupported link type for serialization: $(typeof(m.link))") + end + return Dict( + "coef" => m.coef, + "intercept" => m.intercept, + "link" => linkname, + "hasoffset" => m.hasoffset, + ) +end + +""" + InferenceModel(d::AbstractDict) -> InferenceModel + +Reconstruct an [`InferenceModel`](@ref) from a `Dict` produced by +[`to_dict`](@ref). +""" +function InferenceModel(d::AbstractDict) + linkname = d["link"] + constructor = get(LINK_NAME_TO_CONSTRUCTOR, linkname) do + error("Unsupported link name: $linkname") + end + return InferenceModel(Float64.(d["coef"]), d["intercept"], constructor(), d["hasoffset"]) +end + +""" + predict(m::InferenceModel, newX::AbstractMatrix; offset=eltype(newX)[]) + +Predicted values from an `InferenceModel`. If the original model was fit with +an offset, `offset` must be supplied here with one value per row of `newX` +(mirroring `predict(m::RegularizedModel, newX; offset=...)`). +""" +function StatsBase.predict(m::InferenceModel, newX::AbstractMatrix{T}; + offset::AbstractVector{<:Real}=T[]) where T + X = m.intercept ? [ones(T, size(newX, 1), 1) newX] : newX + eta = X * m.coef + if m.hasoffset + length(offset) == size(newX, 1) || + throw(ArgumentError("model was fit with an offset, so `offset` kwarg must have length size(newX, 1)")) + eta = eta .+ offset + else + isempty(offset) || + throw(ArgumentError("model was fit without an offset, so the `offset` kwarg does not make sense")) + end + linkinv.(m.link, eta) +end + +""" + write_json(io::IO, m) + write_json(path::AbstractString, m) + +Save the coefficients, intercept flag, link, and offset requirement of a +fitted `RegularizedModel` (e.g. a `LassoModel`, or an already-constructed +[`InferenceModel`](@ref)) as JSON, sufficient to later call `predict` via +[`read_json`](@ref). Requires JSON.jl to be loaded. +""" +function write_json end + +""" + read_json(io::IO) -> InferenceModel + read_json(path::AbstractString) -> InferenceModel + +Load a model saved by [`write_json`](@ref) as an [`InferenceModel`](@ref), +which supports `predict`. Requires JSON.jl to be loaded. +""" +function read_json end diff --git a/test/Project.toml b/test/Project.toml index e8ea3ba..3b2cee5 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -4,6 +4,7 @@ DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" GLM = "38e38edf-8417-5370-95a0-9cbb8c7f171a" GLMNet = "8d5ece8b-de18-5317-b113-243142960cc6" +JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MLBase = "f0e99cf1-93fa-52ec-9ecc-5026115318e0" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" @@ -12,4 +13,5 @@ StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [compat] -DataFrames = "1" \ No newline at end of file +DataFrames = "1" +JSON = "1" diff --git a/test/json.jl b/test/json.jl new file mode 100644 index 0000000..896f4d3 --- /dev/null +++ b/test/json.jl @@ -0,0 +1,235 @@ +using JSON + +# extra distributions used only by the JSON serialization tests below, generated +# with a mean-matching random draw (mirrors the `randdist` pattern in lasso.jl) +randdist(::Gamma, x) = rand(testrng, Gamma(2.0, x / 2)) +randdist(::InverseGaussian, x) = rand(testrng, InverseGaussian(x, 1.0)) +randdist(::NegativeBinomial, x) = rand(testrng, NegativeBinomial(10.0, 10.0 / (10.0 + x))) + +""" +Check that `m`'s prediction matches after roundtripping through `InferenceModel`/ +`to_dict` (JSON-independent) and through `write_json`/`read_json` (JSON-based). +""" +function test_json_roundtrip(m, X; offset::AbstractVector{<:Real}=Float64[]) + expected = isempty(offset) ? predict(m, X) : predict(m, X; offset=offset) + + d = Lasso.to_dict(Lasso.InferenceModel(m)) + m_dict = Lasso.InferenceModel(d) + actual_dict = isempty(offset) ? predict(m_dict, X) : predict(m_dict, X; offset=offset) + @test actual_dict ≈ expected + + io = IOBuffer() + write_json(io, m) + m_json = read_json(IOBuffer(take!(io))) + actual_json = isempty(offset) ? predict(m_json, X) : predict(m_json, X; offset=offset) + @test actual_json ≈ expected + + return d +end + +@testset "JSON roundtrip" begin + @testset "$(typeof(dist).name.name) $(typeof(link).name.name)" for (dist, link) in + ((Normal(), IdentityLink()), (Binomial(), LogitLink())) + Random.seed!(testrng, 371) + (X, y) = genrand(Float64, dist, link, 200, 5, false) + + @testset "$(intercept ? "w/" : "w/o") intercept" for intercept in (true, false) + m = fit(LassoModel, X, y, dist, link; intercept=intercept, select=MinAICc()) + + @testset "path" begin + mktempdir() do dir + path = joinpath(dir, "model.json") + write_json(path, m) + m2 = read_json(path) + + @test m2 isa Lasso.InferenceModel + @test predict(m, X) ≈ predict(m2, X) + end + end + + @testset "IO" begin + io = IOBuffer() + write_json(io, m) + m2 = read_json(IOBuffer(take!(io))) + + @test m2 isa Lasso.InferenceModel + @test predict(m, X) ≈ predict(m2, X) + end + + @testset "InferenceModel construction" begin + m2 = Lasso.InferenceModel(m) + @test m2 isa Lasso.InferenceModel + @test predict(m, X) ≈ predict(m2, X) + @test Lasso.InferenceModel(m2) === m2 + end + + @testset "to_dict roundtrip (JSON-independent)" begin + d = Lasso.to_dict(Lasso.InferenceModel(m)) + m2 = Lasso.InferenceModel(d) + + @test m2 isa Lasso.InferenceModel + @test predict(m, X) ≈ predict(m2, X) + end + end + end + + @testset "distribution/link coverage" begin + Random.seed!(testrng, 371) + (X, _) = genrand(Float64, Normal(), IdentityLink(), 200, 5, false) + beta = [(-1)^j * exp(-2 * (j - 1) / 20) for j = 1:5] + eta = 0.3 .* (X * beta) + + # generate y in a domain-valid way *independent* of the link under test + # (fitting with a non-canonical link doesn't require y to have been + # generated through that link's inverse) + positive_mu = exp.(eta) + unit_mu = GLM.linkinv.(LogitLink(), eta) + + @testset "$(typeof(dist).name.name) $(typeof(link).name.name)" for (dist, link, mu) in ( + (Poisson(), LogLink(), positive_mu), + (Gamma(), InverseLink(), positive_mu), + (Gamma(), LogLink(), positive_mu), + (InverseGaussian(), InverseSquareLink(), positive_mu), + (NegativeBinomial(10.0, 0.5), LogLink(), positive_mu), + (Binomial(), ProbitLink(), unit_mu), + (Binomial(), CloglogLink(), unit_mu), + ) + y = [randdist(dist, m) for m in mu] + m = fit(LassoModel, X, Float64.(y), dist, link; select=MinAICc()) + + d = test_json_roundtrip(m, X) + @test d["link"] == string(nameof(typeof(link))) + end + end + + @testset "weights" begin + Random.seed!(testrng, 371) + (X, y) = genrand(Float64, Normal(), IdentityLink(), 200, 5, false) + wts = rand(testrng, 200) .+ 0.5 + + m = fit(LassoModel, X, y; wts=wts, select=MinAICc()) + test_json_roundtrip(m, X) + end + + @testset "penalty_factor" begin + Random.seed!(testrng, 371) + (X, y) = genrand(Float64, Normal(), IdentityLink(), 200, 5, false) + penalty_factor = [0.0, 1.0, 1.0, 0.5, 2.0] + + m = fit(LassoModel, X, y; penalty_factor=penalty_factor, select=MinAICc()) + test_json_roundtrip(m, X) + end + + @testset "sparse X" begin + Random.seed!(testrng, 371) + (X, y) = genrand(Float64, Normal(), IdentityLink(), 200, 5, true) + + m = fit(LassoModel, sparse(X), y; select=MinAICc()) + test_json_roundtrip(m, Matrix(X)) + end + + @testset "segment selectors" begin + Random.seed!(testrng, 371) + (X, y) = genrand(Float64, Normal(), IdentityLink(), 200, 5, false) + path = fit(LassoPath, X, y) + + @testset "$(typeof(select))" for select in + (MinAIC(), MinAICc(), MinBIC(), MinCVmse(path), MinCV1se(path)) + Random.seed!(421) + m = fit(LassoModel, X, y; select=select) + test_json_roundtrip(m, X) + end + end + + @testset "GammaLassoModel" begin + Random.seed!(testrng, 371) + (X, y) = genrand(Float64, Normal(), IdentityLink(), 200, 5, false) + + m = fit(GammaLassoModel, X, y; γ=1.0, select=MinAICc()) + test_json_roundtrip(m, X) + end + + @testset "offset" begin + Random.seed!(testrng, 371) + (X, y) = genrand(Float64, Poisson(), LogLink(), 200, 5, false) + offset = fill(0.1, length(y)) + + m = fit(LassoModel, X, y, Poisson(), LogLink(); offset=offset, select=MinAICc()) + d = test_json_roundtrip(m, X; offset=offset) + @test d["hasoffset"] + + m2 = Lasso.InferenceModel(d) + @test_throws ArgumentError predict(m2, X) + @test_throws ArgumentError predict(m2, X; offset=offset[1:end-1]) + + m0 = fit(LassoModel, X, y, Poisson(), LogLink(); select=MinAICc()) + d0 = Lasso.to_dict(Lasso.InferenceModel(m0)) + @test !d0["hasoffset"] + m0_2 = Lasso.InferenceModel(d0) + @test_throws ArgumentError predict(m0_2, X; offset=offset) + end + + @testset "unsupported model types" begin + Random.seed!(testrng, 371) + (X, y) = genrand(Float64, Normal(), IdentityLink(), 50, 3, false) + path = fit(LassoPath, X, y) + @test_throws ArgumentError Lasso.InferenceModel(path) + + gpath = fit(GammaLassoPath, X, y) + @test_throws ArgumentError Lasso.InferenceModel(gpath) + + fl = fit(FusedLasso, y, 1.0) + @test_throws ArgumentError Lasso.InferenceModel(fl) + + tf = fit(TrendFilter, y, 1, 1.0) + @test_throws ArgumentError Lasso.InferenceModel(tf) + end + + @testset "TableRegressionModel wrapper" begin + data = DataFrame(X1=randn(testrng, 100), X2=randn(testrng, 100)) + data.Y = 1 .+ 2 .* data.X1 .- data.X2 .+ 0.1 .* randn(testrng, 100) + + m = fit(LassoModel, @formula(Y ~ X1 + X2), data; select=MinAICc()) + newX = Matrix(data[:, [:X1, :X2]]) + + m2 = Lasso.InferenceModel(m) + @test predict(m, data) ≈ predict(m2, newX) + + d = Lasso.to_dict(m2) + m2_from_dict = Lasso.InferenceModel(d) + @test predict(m, data) ≈ predict(m2_from_dict, newX) + + mktempdir() do dir + path = joinpath(dir, "model.json") + write_json(path, m) + m3 = read_json(path) + @test predict(m, data) ≈ predict(m3, newX) + end + end + + @testset "serialization stability" begin + # This Dict/JSON is a hand-written, frozen stand-in for a model + # serialized by an older version of Lasso.jl (it is not produced by + # fitting a model here). This test should NOT fail in a non-breaking + # release: its purpose is to make sure models serialized by older + # versions of the package can still be deserialized and used for + # inference. If this test fails, check whether the change is an + # intentional, documented breaking change to the serialization format. + d = Dict( + "coef" => [1.0, 2.0, -1.0], + "intercept" => true, + "link" => "IdentityLink", + "hasoffset" => false, + ) + json_str = """{"coef":[1.0,2.0,-1.0],"intercept":true,"link":"IdentityLink","hasoffset":false}""" + + X = [1.0 2.0; 3.0 -1.0; 0.0 0.0] + expected = [1.0, 8.0, 1.0] + + m_dict = Lasso.InferenceModel(d) + @test predict(m_dict, X) ≈ expected + + m_json = read_json(IOBuffer(json_str)) + @test predict(m_json, X) ≈ expected + end +end diff --git a/test/runtests.jl b/test/runtests.jl index 2f0b528..d33b40f 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -13,5 +13,6 @@ include("fusedlasso.jl") include("trendfiltering.jl") include("cross_validation.jl") include("segselect.jl") +include("json.jl") end