Julia:高效易用的数值计算/优化编程语言

1,148 阅读12分钟
原文链接: zhuanlan.zhihu.com

今天忙里偷闲,简单科普一下Julia这款非常好用的编程语言。个人感觉,对做优化/数值计算的人来说,这是目前最好用的语言了。对我来说,3年来自从用了Julia我就再也不用Matlab了,连Python我觉得也基本没有用的需求。

The Julia Language

一、Julia中的数值计算入门

首先,Julia的LLVM-based just-in-time (JIT)编译器在提供较低难度的编写环境(和C相比)的情况下,仍能在数值计算计算时间的表现上和C达到类似的运算速度(见下图,详见Julia Micro-Benchmarks)。测试的operation包括 \pi 的级数求和,递归求斐波那契和/快速排序,输出到文件,矩阵乘法/特征值等。


例子1:简单随机游走(Simple Random Walk)的hitting time

walk(L)函数定义了一个从 x=2 为起点的,右边界为 x=L ,左边界为 x=1 的简单随机游走(左移右移一格的概率永远都是0.5)。在这个例子中,我们主要考虑到达左边界的hitting time,times(L,N)函数即是返回N次walk(L)第一次达到左边界的时间,repeat(L,N,M)函数让我们可以重复多次walk(即蒙特卡洛模拟),并返回平均的hitting time。具体来说,一共M次模拟,每次模拟重复N次取平均。

function walk(L=50)
    t = 2
    x = 2
    while x != 1
        if rand() < 0.5
            x += 1
        else
            x -= 1
        end 
        if x > L
            x = L
        end
        t += 1
    end
    return t
end

function run(L, N)
    times = Int64[]
    for i in 1:N
        push!(times, walk(L))
    end
    return times
end

function repeat(L, N, M)
    [mean(run(L, N)) for i in 1:M]
end

我们的蒙特卡洛模拟返回的平均hitting time是99.90814250000001,并输出了相应的pmf的histogram。

data = repeat(50, 1000, 10000);
using Plots
mean(data)
histogram(data)

那么我们可以看到pmf是比较类似高斯分布的,且mean集中在 2L 左右。当然这是显然的,因为其实我们可以手算出来进行验证。定义 \mu_i=\mathbb{E}[\min\{ n\geq 0| X_n = 1 \}|X_0 = i ] .那么容易通过 \mu_1=2,\mu_2=1+\frac{1}{2}\mu_1+\frac{1}{2}\mu_3,\mu_3=1+\frac{1}{2}\mu_2+\frac{1}{2}\mu_4,\ldots,\mu_{L-1}=1+\frac{1}{2}\mu_{L-2}+\frac{1}{2}\mu_L,\mu_L=1+\frac{1}{2}\mu_{L-1}+\frac{1}{2}\mu_L 得到 \mu_2=2L. 至于分布看起来像高斯是因为中心极限定理。

因为这个random walk也可以看成一个马尔可夫链,所以我们也可以通过状态转移矩阵的乘积去计算它的stationary distribution,好吧这里主要是为了展示更多Julia的命令和比较一下array comprehension和(稀疏)矩阵乘法的速度。我们利用BenchmarkTools的package比较这三种算法的速度。让 L=500 且迭代 1000 次(tmax=1000)。

# Array comprehensions
function evolve(p)
    L = length(p)
    
    p_new = vec(zeros(L,1))
    p_new[1] = p[1] + p[2]/2
    p_new[2] = p[3] / 2
    
    for i in 3:L-1
        p_new[i] = (p[i-1] + p[i+1]) / 2
    end
    
    p_new[L] = (p[L-1] + p[L]) / 2
    
    return p_new
end

@benchmark begin
    p = [i == 2 ? 1//1 : 0//1 for i in 1:L];
    P = zeros(t_max,L)
    P[1,:] = evolve(p)
    for i = 2:t_max
        P[i,:] = evolve(P[i-1,:])
    end
end
# Dense matrix
@benchmark begin
    p = [i == 2 ? 1//1 : 0//1 for i in 1:L];
    P = zeros(t_max,L)
    Trans = zeros(L,L)
    Trans[1,1] = 1
    Trans[1,2] = 0.5
    Trans[2,3] = 0.5
    for i = 3:L-1
        Trans[i,i-1] = 0.5
        Trans[i,i+1] = 0.5
    end
    Trans[L,L-1] = 0.5
    Trans[L,L] = 0.5
    P[1,:] = Trans * vec(p)
    for i = 2:t_max
        P[i,:] = Trans * vec(P[i-1,:])
    end
end
# Sparse matrix
@benchmark begin
    t_max = 1000
    p = [i == 2 ? 1//1 : 0//1 for i in 1:L];
    P = zeros(t_max,L)
    dl = 0.5 * vec(ones(L-1,1))
    dl[1] = 0
    d = vec(zeros(L,1))
    d[1] = 1
    d[L] = 0.5
    du = 0.5 * vec(ones(L-1,1))
    Trans = Tridiagonal(dl,d,du)
    P[1,:] = Trans * vec(p)
    for i = 2:t_max
        P[i,:] = Trans * vec(P[i-1,:])
    end
end

注意到在稀疏矩阵那里我们利用了三对角矩阵(Tridiagonal)的数据结构,而benchmarkTools都做了几百次的sample然后输出该代码块运算时间的各种统计量。我们于是知道稀疏矩阵快于array comprehension快于naive的稠密矩阵。在数值计算中,使用最有效的数据结构和代码命令永远是很重要的。至于如果想验证(sanity check)这三个代码块是否给出了一致的输出结果,可以使用@test_approx_eq命令,这里不多赘述。


例子2:稀疏矩阵,自定义数据结构,分布式数据结构(初步)

我们在上例中看到了一类特殊的稀疏矩阵——三对角矩阵类型的应用。这里我们指出更一般的稀疏矩阵类型在Julia中名为sparse()。一个简单例子如下:

G = sparse(vec([1;2;3;1;1;2;3]),vec([1;2;3;2;3;1;1]),vec([1.;3;5;2;2;2;2]))

我们看到这种类型的矩阵便只会存储非0的entry和其对应的值,这和Matlab等都是很类似的。接下来我们尝试自定义新的矩阵类型(数据结构),我们叫他们SymArrow矩阵,一类对称矩阵,除了对角线和第一行/列其它entry都是0的稀疏矩阵。

type SymArrowFloat
    Diag::Array{Float64,1}
    Row::Array{Float64,1}
end

import Base: full 
Base.full(A::SymArrowFloat) = full(sparse(round(Int64,[1:size(A.Diag,1);vec(ones(size(A.Row)));2:size(A.Diag,1)]),round(Int64,[1:size(A.Diag,1);2:size(A.Diag,1);vec(ones(size(A.Row)))]),[A.Diag;A.Row;A.Row] ))

import Base: promote_rule, convert, show
Base.show(io::IO, A::SymArrowFloat)=print(io,sparse(round(Int64,[1:size(A.Diag,1);vec(ones(size(A.Row)));2:size(A.Diag,1)]),round(Int64,[1:size(A.Diag,1);2:size(A.Diag,1);vec(ones(size(A.Row)))]),[A.Diag;A.Row;A.Row] ) )

import Base.+
+(A::SymArrowFloat,B::SymArrowFloat) = SymArrowFloat(A.Diag+B.Diag, A.Row+B.Row)

import Base.*
*(A::SymArrowFloat,v::Array{Float64,1}) = [dot(vec([A.Diag[1];A.Row]),vec(v));[v[1]*A.Row[i-1]+v[i]*A.Diag[i] for i=2:size(vec(v),1)]]

我们给这类新的数据结构自定义了full()和show()函数,前者是将一个稀疏矩阵变成稠密矩阵,后者则是显示该数据类型。同样,我们也定义了该数据类型的加减和乘法。

# Function full works
G = SymArrowFloat([1.;3;5],[2;2.])
full(G)
# Function show works
G
# + works
G+G
# * works
G*rand(3)

这里不贴运行结果了,大家可以自行尝试都是work的。这边我们指出,如果想specify有理数,复数为值域,也可以单独定义(其实主要只是为了展示Julia的数域类型233)。

type SymArrow{T}
    Diag::Array{T,1}
    Row::Array{T,1} 
end

import Base: full 
Base.full{T}(A::SymArrow{T}) = full(sparse(round(Int64,[1:size(A.Diag,1);vec(ones(size(A.Row)));2:size(A.Diag,1)]),round(Int64,[1:size(A.Diag,1);2:size(A.Diag,1);vec(ones(size(A.Row)))]),[A.Diag;A.Row;A.Row] ))

import Base.+
+{T}(A::SymArrow{T},B::SymArrow{T}) = SymArrow{T}(A.Diag+B.Diag, A.Row+B.Row)

import Base.*
*{T}(A::SymArrow{T},v::Array{T,1}) = [dot(vec([A.Diag[1];A.Row]),vec(v));[v[1]*A.Row[i-1]+v[i]*A.Diag[i] for i=2:size(vec(v),1)]]

import Base: promote_rule, convert, show
Base.show{T}(io::IO, A::SymArrow{T})=print(io,sparse(round(Int64,[1:size(A.Diag,1);vec(ones(size(A.Row)));2:size(A.Diag,1)]),round(Int64,[1:size(A.Diag,1);2:size(A.Diag,1);vec(ones(size(A.Row)))]),[A.Diag;A.Row;A.Row] ) )

我们看到和前面相比我们就是特地把{T}加到了和Array相关的Type定义里。这里T可以是整数,有理数,复数,BigFloat等各种数域类型。

G = SymArrow{Rational}([1.;3;5],[2;2.])
G = SymArrow{Complex}([1.+0im;3;5],[2;2.-0im])
G = SymArrow{BigFloat}([1.;3;5],[2;2.])
G*[BigFloat(1);2;3]

如果想省事的话,我们也可以用AbstractMatrix类型,不过这种情况下的运算速度会稍微不如之前的定义方式(因为这时候无法利用稀疏性。大家可以自行验证,之前的那些操作在这里都可以自动work)。

type SymArrow2{T} <: AbstractMatrix{T}
    Diag::Array{T,1}
    Row::Array{T,1}  
end

import Base: size, getindex

size{T}(A::SymArrow2{T}) = (length(A.Diag), length(A.Diag))

function getindex{T}(A::SymArrow2{T}, i, j)  
    if i == j
        return A.Diag[i]
    elseif i == 1
        return A.Row[j-1]
    elseif j == 1
        return A.Row[i-1]
    else
        return zero(T)  # otherwise return zero of type T
    end
end

这个例子最后我们看一下Julia支持的分布式数据结构(这里只简单介绍array),和相关的一些分布式运算。至于真正要用Julia做并行计算不在这篇文章里重点谈,可能以后可以再专门科普。我们这里只涉及DistributedArrays.jl pacakge。例子是对一个array分布式地累加求和(cumulative sum)。我们用4个核来分布式地完成操作。

using DistributedArrays
nprocs()  ==  1  &&  addprocs(4)
xd = @DArray [i for i = 0:39]
[DistributedArrays.chunk(xd, i) for i = 1:4]

我们可以看到array被自动partition成4部分分配到了相应的区域里。为了并行运算,我们需要定义prefix()函数,主要是map-reduce框架里map的部分。

@everywhere  function prefix!(op, v, v0 = 0)
    v[1] += v0
    for i = 2:length(v)
        v[i] = op(v[i-1], v[i])
    end
    return v
end

@everywhere function prefixlag!(op, v, v0 = 0)
    vi = v[1]
    v[1] = v0
    for i = 2:length(v)
        tmp  = op(v[i - 1], vi)
        vi   = v[i]
        v[i] = tmp
    end
    return v
end

我们于是便可以分布式地得到cumsum。

yd1 = map_localparts(t -> [sum(t)], xd)
yd2 = prefixlag!(+, Array(yd1))
yd3 = map_localparts((t,s) -> prefix!(+, t, s[1]), xd, distribute(yd2))

例子3:奇异值分解(SVD)

这个例子是我这个回答里的源码。覃含章:SVD 降维体现在什么地方?

using Images, Colors, ImageCore, ImageView
img = load("Albert Einstein-rare-pics56.jpg")
img_RGB = float(Array(channelview(img)))
R_pixel = img_RGB[1,:,:];
G_pixel = img_RGB[2,:,:];
B_pixel = img_RGB[3,:,:];
U1,S1,V1 = svd(R_pixel);
U2,S2,V2 = svd(G_pixel);
U3,S3,V3 = svd(B_pixel);
U1[:,1:100]*diagm(S1[1:100])*V1[1:100,:]

# Full SVD: perfect reconstruction
Images.colorim(cat(3,(U1*diagm(S1)*V1')',(U2*diagm(S2)*V2')',(U3*diagm(S3)*V3')'))
# 25% SVD reconstruction: near perfect
S1[101:400] .= 0
S2[101:400] .= 0
S3[101:400] .= 0
Images.colorim(cat(3,(U1*diagm(S1)*V1')',(U2*diagm(S2)*V2')',(U3*diagm(S3)*V3')'))

# 2.5% SVD reconstruction: obviously blurred
S1[11:400] .= 0
S2[11:400] .= 0
S3[11:400] .= 0
Images.colorim(cat(3,(U1*diagm(S1)*V1')',(U2*diagm(S2)*V2')',(U3*diagm(S3)*V3')'))

# 1% SVD reconstruction: heavily blurred
S1[5:400] .= 0
S2[5:400] .= 0
S3[5:400] .= 0
Images.colorim(cat(3,(U1*diagm(S1)*V1')',(U2*diagm(S2)*V2')',(U3*diagm(S3)*V3')'))

# only 2 nonzero entries: stripes
S1[3:400] .= 0
S2[3:400] .= 0
S3[3:400] .= 0
Images.colorim(cat(3,(U1*diagm(S1)*V1')',(U2*diagm(S2)*V2')',(U3*diagm(S3)*V3')'))

二、JuMP入门:线性规划,锥规划,混合整数规划

除了数值计算,Julia语言对于优化工作者很方便的一点是,类似于Matlab中的YALMIP(yalmip.github.io/),提供了high level的优化模型的输入接口,并调用不同的solver进行优化模型的求解。在Julia中,这主要是通过JuMP这个package来实现的。

目前支持的solver一览:(详见Optimization packages for the Julia language

在建模速度上,严格好于其它市面上的AML(algebraic modeling language)接口,且不输于商业AML!

来源:Dunning, Iain, Joey Huchette, and Miles Lubin. &amp;quot;JuMP: A modeling language for mathematical optimization.&amp;quot; SIAM Review 59.2 (2017): 295-320.

例子1:支持向量机(Support Vector Machine, SVM)【线性规划】

这个例子里,我给出SVM利用线性规划(linear programming)在Julia中的实现方式,linear kernel和RBF kernel两种kernel function。具体来说,这里考虑的是soft-margin的SVM的对偶表示,相应对slack variable的cost是 C.

using PyPlot, JuMP

function linear_K(X)
    K = zeros(size(X,1),size(X,1))
    for i = 1:size(X,1)
        for j = 1:size(X,1)
            K[i,j] = dot(vec(X[i,:]), vec(X[j,:]))
        end
    end
    return K
end

function RBF_K(X,γ)
    K = zeros(size(X,1),size(X,1))
    for i = 1:size(X,1)
        for j = 1:size(X,1)
            K[i,j] = exp( -γ*sum(( X[i,:]-X[j,:] ).^2) )
        end
    end
    return K
end

function SVM_K(X,Y,C,K_type,γ=1)
    n = length(Y)
    M_dual = Model(solver = GurobiSolver(OutputFlag=0))
    @variable(M_dual, 0<=α[1:n]<=C)
    @constraint(M_dual, sum(Y.*α)==0)
    if K_type == "linear"
        @objective(M_dual, Max, -0.5*sum(α'*linear_K(X)*α) + sum(α))
    elseif K_type == "RBF"
        @objective(M_dual, Max, -0.5*sum(α'*RBF_K(X,γ)*α) + sum(α))
    else
        error("Wrong kernel type!")
    end
    solve(M_dual)
    # return optimal α and margin
    return getvalue(α), 1/(getobjectivevalue(M_dual)-sum(getvalue(α)))*(-2)
end

注意到在这里的JuMP model里我调用了Gurobi作为solver。


例子2:计算图的stable number【半正定规划】

给定一个无向图G=(V,E),一个stable set指的是节点V的一个子集,其中没有任何两个节点是相连的。stable number指的就是G上所有stable set里最大的节点数量。这个问题可以relax成一个半正定规划(semidefinite programming)问题:

\begin{align} \alpha(G) = & \max_X~ \text{tr}(ee^TX) \\ &\text{s.t. } X_{ij} = 0,~ \text{if } (i,j)\in E,\\ & \text{tr}(X) = 1,\\ & X \succeq 0. \end{align}

Julia代码如下。核心就是利用@SDconstraint这个macro,以及注意这里我调用了Mosek作为solver。这里有一些具体算例:Copositive Programming简介

using JuMP, Mosek
function StableSDP(A)
    n = size(A,1)
    StableDD = Model(solver = MosekSolver(MSK_IPAR_LOG = 0))
    @variable(StableDD, X[1:n,1:n])
    @SDconstraint(StableDD, X>=0)
    @constraint(StableDD, X.>=0)
    @constraint(StableDD, sum(sum(A.*X)) == 0 )
    @constraint(StableDD, sum(X[i,i] for i=1:n) == 1)
    @objective(StableDD, Max, sum(sum(X)))
    solve(StableDD)
    return getobjectivevalue(StableDD)
end

例子3:Gurobi进阶【混合整数规划】

在这里,我们指出使用JuMP的接口我们仍能具有比较自由地和Gurobi进行互动的自由度。比如说,我们可以加入lazy constraint(即一开始只specify一部分constraint,随着求解过程慢慢加入其它constraint)或者user cut,这都可以穿插在branch and bound过程中。以及Gurobi也支持warmstart,即user可以自己提供一个初始值(比如用某些heuristic得到的,加速branch and bound);这方面的详细教程可看 Solver Callbacks

同样,也可以利用Gurobi做column generation(和lazy constraint类似,只不过一开始只加入一部分variable,随着求解过程慢慢加入其它variable)。然而不幸的是,Gurobi一直还没有支持branch and price(即在整数规划中做column generation),因此目前CG只能用在连续优化问题当中。这方面,Chiwei Yan- JuliaTutorial 的cutting stock problem教程是很好的学习资料。


三、交互性编程

我们指出,类似于IPython,Julia中的IJulia package也可以让Julia的所有编译过程在Jupyter notebook里进行。我个人是很喜欢在这个环境里进行Julia编程的(详见:JuliaLang/IJulia.jl)。当然,我知道也有不少人喜欢Juno的:Juno, the Interactive Development Environment 可能这个更有码农的感觉吧hhh

进一步的,我们可以有很多交互式的操作,这在Julia中主要通过Interact package实现。比如,我们可以自定义slider,按钮等对一个参数曲线进行互动。一个例子见如下视频,或者JuliaGizmos/Interact.jl

Interact in Julia

至于各种画图,我倾向于使用Plots这个package。入门可以见:Plots - powerful convenience for visualization in Julia.


四、写在最后

自然,本文给出的只是很少的一些例子和对Julia这门编程语言的最基本的介绍。无论你只是希望有个方便的语言调用solver,或者做数值计算居多,还是比较高级的优化算法专家,多实践,一边"get hands dirty"一边学习我觉得总是最有效率的。

首当其冲的是Julia官网上提供的大量学习资源:包括视频,具体的算例,和各种pdf教程。

Learning Julia

作为Julia cofounder之一的Prof. Edelman的数值计算课程的个人主页也是非常不错的参考资料:(主要是那个Github链接)

Modern Numerical Computing with Julia

最后提一点个人的小期许。熟悉Julia的同学们也可以在知乎上发布相关的教程,比如一些更进阶的应用(鲁棒优化,深度学习,符号计算,大规模分布式计算等),我这边的个人专栏和『运筹OR帷幄』大数据人工智能时代的运筹学 专栏也都可以帮忙推广发布。希望可以更壮大一点Julia和运筹学/优化/OR这个社群呢~


相关的知乎回答:

Julia解决了C++/Python/Matlab的哪些痛点? 怎么看待新出的 Julia 语言?它的前景怎么样?