add serialization features to be able to robustly load and save models

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"

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

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",
     ],
     )

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
+```

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

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")
 

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

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"
+DataFrames = "1"
+JSON = "1"