-
Notifications
You must be signed in to change notification settings - Fork 7
/
geom_smooth.jl
129 lines (103 loc) · 3.21 KB
/
geom_smooth.jl
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
"""
geom_smooth(aes(...), ...)
geom_smooth(plot::GGPlot, aes(...), ...)
Represent data as a smoothed or linear fit.
# Arguments
- `plot::GGPlot` (optional): a plot object to add this geom to
- `aes(...)`: the names of the columns in the DataFrame that will be used in the mapping
- `...`: options that are not mapped to a column (passed to Makie.Lines)
# Required Aesthetics
- `x`
- `y`
# Optional Aesthetics (see [`aes`](@ref))
- NA
# Optional Arguments
- `method`: either "smooth" (default, loess fit) or "lm" (linear fit)
- `color` / `colour`
- `linewidth`
- `alpha`
- `linestyle` / `linetype`
# Examples
```julia
xs = range(0, 2pi, length=30)
ys = sin.(xs) .+ randn(length(xs)) * 0.5
df = DataFrame(x = xs, y = ys)
ggplot(df, @aes(x = x, y = y)) + geom_smooth() + geom_point()
ggplot(penguins, @aes(x = bill_length_mm, y = bill_depth_mm)) +
geom_smooth(color=:red, linewidth=10, alpha=0.5)
```
"""
function geom_smooth(plot::GGPlot, args...; kwargs...)
return plot + geom_smooth(args...; kwargs...)
end
function geom_smooth(args...; kwargs...)
aes_dict, args_dict, transforms = extract_aes(args, kwargs)
args_dict["geom_name"] = "geom_smooth"
if get(args_dict, "method", "smooth") == "lm"
return [build_geom(aes_dict,
args_dict,
["x", "y"],
:Lines,
stat_linear,
transforms),
build_geom(aes_dict,
args_dict,
["x", "lower", "upper"],
:Band,
stat_linear,
transforms)]
end
return build_geom(aes_dict,
args_dict,
["x", "y"],
:Lines,
stat_loess,
transforms)
end
function stat_loess(
aes_dict,
args_dict,
required_aes::Vector{String},
plot_data::DataFrame
)
x = plot_data[!, aes_dict["x"]]
y = plot_data[!, aes_dict["y"]]
model = Loess.loess(x, y; span = .75, degree = 2)
x̂ = range(extrema(x)..., length=200)
ŷ = Loess.predict(model, x̂)
return_data = DataFrame(
String(aes_dict["x"]) => x̂,
String(aes_dict["y"]) => ŷ
)
return (aes_dict, args_dict, required_aes, return_data)
end
function stat_linear(
aes_dict,
args_dict,
required_aes::Vector{String},
plot_data::DataFrame
)
x = plot_data[!, aes_dict["x"]]
y = plot_data[!, aes_dict["y"]]
# thanks AlgebraOfGraphics
function add_intercept_column(x::AbstractVector{T}) where {T}
mat = similar(x, float(T), (length(x), 2))
fill!(view(mat, :, 1), 1)
copyto!(view(mat, :, 2), x)
return mat
end
lin_model = GLM.lm(add_intercept_column(x), y)
x̂ = range(extrema(x)..., length = 100)
pred = DataFrame(
GLM.predict(lin_model, add_intercept_column(x̂); interval = :confidence)
)
aes_dict["upper"] = :upper
aes_dict["lower"] = :lower
return_data = DataFrame(
String(aes_dict["x"]) => x̂,
String(aes_dict["y"]) => pred.prediction,
"lower" => pred.lower,
"upper" => pred.upper
)
return (aes_dict, args_dict, required_aes, return_data)
end