Julia: типы, мультиметоды и арифметика над полиномами

В этой публикации речь пойдёт об основной, на мой взгляд, отличительной особенности языка Julia — представлении функций в виде методов с множественной диспетчеризацией. Это позволяет повысить производительность вычислений, не снижая читаемости кода и не портя абстрагируемость, с одной стороны, и позволяет работать с математическими понятиями в более привычной нотации, с другой. Для примера рассмотрен вопрос единообразной (с точки зрения линейных операций) работы с полиномами в представлении списка коэффициентов и с интерполяционными полиномами.

Базовый синтаксис

Краткое введение для тех, кто не в курсе. Julia — как-бы-скриптовый язык, имеет REPL (read-evaluate-print loop, т.е. интерактивную оболочку). С первого взгляда выглядит довольно похоже на, например, Python или MATLAB.

Арифметические операции

Арифметика примерно такая же, как и везде: +, -, *, /, ^ для возведения в степень и т.д.
Сравнение: >, <, >=, <=, ==, != и т.д.
Присваивание: =.
Особенности: деление через / всегда возвращает дробное число; если нужна целая часть от деления двух целых чисел, нужно пользоваться операцией div(m, n) или инфиксным эквивалентом m ÷ n.

Типы

Числовые типы:

• Целые (Int) — 2, 3, -42
• Беззнаковые целые (UInt) — 0x12345
• С плавающей точкой (Float32, Float64) — 1.0, 3.1415, -Inf, NaN
• Рациональные (Rational) — 3//3, 7//2
• Действительные (Real) — всё вышеперечисленное
• Комплексные (Complex) — 3+4*im, 2//3+2//3*im, 3.0+0.0*im (im — мнимая единица, комплексным считается только число с явно выписанной мнимой частью)
• Число (Number) — всё вышеперечисленное

Строки и символы:

• 'a' — символ (Char)
• "a" — строка (String)

NB: строки, как сейчас во многих языках, иммутабельные.
NB: строки (а также имена переменных) поддерживают Юникод, в том числе и эмодзи.

Массивы:

• x = [1, 2, 3] — задание массива прямым перечислением элементов
• специальные конструкторы: zeros(length) для массива из нулей, ones(length) для массива из единиц, rand(length) для массива из случайных чисел и др.
• поддержка многомерных массивов
• поддержка операций линейной алгебры (сложение массивов, умножение на скаляр, умножение матрицы на вектор и многое другое) в стандартной библиотеке

NB: индексация всех коллекций идёт начиная с единицы.
NB: т.к. язык предназначен для вычислительных задач, массивы — один из наиболее важных типов, к принципам их работы ещё не раз придётся возвращаться.

Кортежи (упорядоченный набор элементов, иммутабельные):

• (2, 5.3, "k") — обычный кортеж
• (a = 3, b = 4) — именованный кортеж

NB: к полям именованного кортежа можно обращаться как по имени через точку, так и по индексу через []

julia> x = (a = 5, b = 12)
(a = 5, b = 12)
julia> x[1]
5
julia> sqrt(x.a^2 + x[2]^2)
13.0

Словари:

julia> x = Dict('a' => 5, 'b' => 12)
Dict{Char,Int64} with 2 entries:
  'a' => 5
  'b' => 12
julia> x['c'] = 13
13
julia> x
Dict{Char,Int64} with 3 entries:
  'a' => 5
  'c' => 13
  'b' => 12

Основные управляющие конструкции языка

1. Переменные автоматически создаются при присваивании. Тип указывать необязательно.

julia> x = 7; x + 2
9
julia> x = 42.0; x * 4
168.0

2. Блок условного перехода начинается с выражения if <condition> и заканчивается словом end. Можно также иметь else-ветку или elseif-ветки:

if x > y
    println("X is more than Y")
elseif x == y
    println("X and Y are equal")
else
    println("X is less than Y")
end

3. Есть две конструкции циклов: while и for. Второй работает как в Python, т.е. проводит итерирование по коллекции. Частое применение — итерирование по диапазону значений, синтаксис которого start[:increment]:end. В отличие от Python, диапазон включает как начальное, так и конечное значения, т.е. пустой диапазон будет не 1:1 (это диапазон из одного значения 1), а 1:0. Конец тела цикла маркируется словом end.

julia> for i in 1:3; print(i, " "); end # диапазон от 1 до 3 с шагом 1 (по умолчанию)
1 2 3 
julia> for i in 1:2:3; print(i, " "); end # диапазон от 1 до 3 с шагом 2
1 3 

4. Функции задаются ключевым словом function, определение функции также завершается словом end. Поддерживаются аргументы со значениями по умолчанию и именованные аргументы.

function square(x)
    return x * x
end

function cube(x)
    x * square(x) # последнее вычисленное значение возвращается из блока кода; return не обязателен
end

function root(x, degree = 2)
    # аргумент degree имеет значение по умолчанию
    return x^(1.0/degree)
end

function greeting(name; times = 42, greet = "hello")
    # именованные аргументы отделяются точкой с запятой
    println(times, " times ", greet, " to ", name)
end

julia> greeting("John")
42 times hello to John
julia> greeting("Mike", greet = "wassup", times = 100500) # именованные аргументы при вызове функции могут стоять в любом порядке
100500 times wassup to Mike

В целом, это всё довольно похоже на Python, если не считать мелких отличий в синтаксисе и того, что блоки кода выделяются не пробелами, а всё-таки ключевыми словами. В простых случаях программы на Python даже транслируются в Julia практически один к одному.
Но есть существенное отличие в том, что в Julia для переменных можно явно указывать типы, что позволяет компилировать программы, получая быстрый код.
Второе существенное отличие — в том, что в Python реализована «классическая» модель ООП с классами и методами, а в Julia реализована модель множественной диспетчеризации.

Аннотации типов и множественная диспетчеризация

Посмотрим, что представляет собой какая-нибудь встроенная функция:

julia> sqrt
sqrt (generic function with 19 methods)

Как показывает нам REPL, sqrt — это обобщённая функция с 19 методами. Что за обобщённая функция и что за методы?

А означает это то, что есть несколько функций sqrt, которые применяются к разным типам аргументов и, соответственно, вычисляют квадратный корень по различным алгоритмам. Посмотреть, какие есть варианты, можно, набрав

julia> methods(sqrt)

Видно, что функция определена для разных типов чисел, а также для матриц.

В отличие от «классического» ООП, где конкретная реализация метода определяется только вызывающим классом (диспетчеризация по первому аргументу), в Julia выбор функции определяется типами (и количеством) всех её аргументов.
При вызове функции с конкретными аргументами из всех её методов выбирается тот, который наиболее точно описывает конкретный набор типов, с которыми функция вызвана, и именно он применяется.

Отличительной особенностью является то, что применяется подход, называемый авторами языка «just ahead-of-time» компиляцией. Т.е. функции компилируются для заданных типов данных при первом вызове, после чего следующие вызовы выполняются гораздо быстрее. Разница между первым и последующими вызовами может быть весьма существенной:

julia> @time sqrt(8) # макрос @time - простое встроенное средство измерения производительности
  0.006811 seconds (3.15 k allocations: 168.516 KiB) # на самом деле, это время и выделение памяти при компиляции
2.8284271247461903
julia> @time sqrt(15)
  0.000002 seconds (5 allocations: 176 bytes) # 5 выделений памяти - это от вызова макроса @time
3.872983346207417

В плохом случае каждый вызов функции — это проверка типов получаемых аргументов и поиск нужного метода в списке. Однако, если компилятору давать подсказки, проверки можно исключить, что приведёт к более быстрому коду.

Для примера, рассмотрим вычисление суммы

function mysqrt(num)
# если аргумент положителен - возвращает обычный квадратный корень
# если нет - преобразует аргумент к комплексному числу и извлекает корень из него
    if num >= 0
        return sqrt(num)
    else
        return sqrt(complex(num))
    end
end

function S(n)
# оставим автоопределение типа
    sum = 0
    sgn = -1
    for k = 1:n
        sum += mysqrt(sgn)
        sgn = -sgn
    end
    return sum
end

function S_typed(n::Integer)
# т.к. уже первое слагаемое получается комплексное, то ответ должен быть комплексным
# тип переменных указывается через 
    sum::Complex = 0.0
    sgn::Int = -1
    for k = 1:n
        sum += mysqrt(sgn)
        sgn = -sgn
    end
    return sum
end

Бенчмарк показывает, что функция S_typed() не только выполняется быстрее, но ещё и не требует выделений памяти при каждом вызове, в отличие от S(). Проблема тут в том, что тип возвращаемого из mysqrt() значения не определён, как и тип правой части выражения

sum = sum + mysqrt(sgn)

Как следствие, компилятор даже не может понять, какого типа будет sum на каждой итерации. А значит, боксинг (прицепление метки типа) переменной и выделение памяти.
Для функции S_typed() компилятор заранее знает, что sum — это комплексное значение, поэтому код получается более оптимизированным (в частности, вызов mysqrt() можно эффективно заинлайнить, приводя возвращаемое значение всегда к Complex).

Что ещё важнее, для S_typed() компилятор знает, что возвращаемое значение имеет тип Complex, а вот для S() тип выходного значения опять не определён, что будет замедлять и все функции, где S() будет вызываться.
Проверить, что компилятор думает о типах, возвращаемых из выражения, можно с помощью макроса @code_warntype:

julia> @code_warntype S(3)
Body::Any # компилятор не знает до вычисления, какого типа будет возвращаемое значение
...
julia> @code_warntype S_typed(3)
Body::Complex{Float64} # компилятор сразу знает возвращаемый тип
...

Если у вас где-то в цикле вызывается функция, для которой @code_warntype не может вывести возвращаемый тип, или для которой он в теле где-то показывает получение значения типа Any — то оптимизация этих вызовов с большой вероятностью даст очень ощутимый прирост производительности.

Составные типы

Программист может определить составные типы данных для своих нужд с помощью конструкции struct:

julia> struct GenericStruct
# внутри блока struct идёт перечисление полей
    name
    b::Int
    c::Char
    v::Vector
end
# конструктор по умолчанию принимает позиционные аргументы 
# и присваивает их полям в том порядке, в котором они идут в объявлении типа
julia> s = GenericStruct("Name", 1, 'z', [3., 0])
GenericStruct("Name", 1, 'z', [3.0, 0.0])

julia> s.name, s.b, s.c, s.v
("Name", 1, 'z', [3.0, 0.0])

Структуры в Julia иммутабельны, т.е., создав экземпляр структуры, поменять значения полей уже нельзя (точнее, нельзя поменять адрес полей в памяти — элементы мутабельных полей, как, например, s.v в примере выше, могут быть изменены). Мутабельные структуры создаются конструкцией mutable struct, синтаксис которой такой же, как и для обычных структур.

Наследование структур в «классическом» смысле не поддерживается, однако есть возможность «наследования» поведения путём объединения составных типов в надтипы или, как они называются в Julia, абстрактные типы. Отношения типов выражаются как A<:B (A — подтип B) и A>:B (A — надтип B). Выглядит примерно так:

abstract type NDimPoint end # абстрактный тип - нужен только как интерфейс

# считаем, что производные типы - это просто кортежи из N чисел
struct PointScalar<:NDimPoint
    x1::Real
end

struct Point2D<:NDimPoint
    x1::Real
    x2::Real
end

struct Point3D<:NDimPoint
    x1::Real
    x2::Real
    x3::Real
end

# документация пишется перед определением функции; поддерживается форматирование Markdown
"""
    mag(p::NDimPoint)

Calculate the magnitude of the radius vector of an N-dimensional point `p`
"""
function mag(p::NDimPoint)
    sqrmag = 0.0
    # т.к. размерность точно неизвестна, нужно итерировать по всем полям структуры
    # имена полей для типа T получаются вызовом fieldnames(T)
    for name in fieldnames(typeof(p))
        sqrmag += getfield(p, name)^2
    end
    return sqrt(sqrmag)
end

"""
    add(p1::T, p2::T) where T<:NDimPoint

Calculate the sum of the radius vectors of two N-dimensional points `p1` and `p2`
"""
function add(p1::T, p2::T) where T<:NDimPoint
    # сложение - уже сложнее, т.к. оба аргумента должны быть одинаковых типов
    # для получения компонентов используется list comprehension
    sumvector = [Float64(getfield(p1, name) + getfield(p1, name)) for name in fieldnames(T)]
    # возвращаем точку того же типа, что и аргументы
    # оператор ... разбивает коллекцию на отдельные аргументы функции, т.е.
    # f([1, 2, 3]...) - это то же, что f(1, 2, 3)
    return T(sumvector...)
end

Case study: Полиномы

Система типов вкупе с множественной диспетчеризацией удобна для выражения математических понятий. Рассмотрим на примере простой библиотеки для работы с полиномами.
Введём два типа полиномов: «канонический», задаваемый через коэффициентами при степенях, и «интерполяционный», задаваемый набором пар (x, f(x)). Для простоты рассматривать будем только действительные аргументы.

Для хранения многочлена в обычной записи подходит структура, имеющая в качестве поля массив или кортеж коэффициентов. Чтобы было совсем иммутабельно, пусть будет кортеж. Таким образом, код для задания абстрактного типа, структуры многочлена и вычисления значения многочлена в заданной точке довольно простой:

abstract type AbstractPolynomial end

"""
    Polynomial <: AbstractPolynomial

Polynomials written in the canonical form
"""
struct Polynomial<:AbstractPolynomial
    degree::Int
    coeff::NTuple{N, Float64} where N # NTuple{N, Type} - тип кортежа из N элементов одинакового типа
end

"""
    evpoly(p::Polynomial, z::Real)

Evaluate polynomial `p` at `z` using the Horner's rule
"""
function evpoly(p::Polynomial, z::Real)
    ans = p.coeff[end]
    for idx = p.degree:-1:1
        ans = p.coeff[idx] + z * ans
    end
    return ans
end

Для интерполяционных полиномов нужна другая структура представления и способ вычисления. В частности, если набор точек интерполяции известен заранее, и один и тот же многочлен планируется вычислять в разных точках, удобна интерполяционная формула Ньютона:

где n_k(x) — базисные полиномы, n₀(x) и для k>0

где x_i — узлы интерполяции.

Из приведённых формул видно, что хранение удобно организовать в виде набора узлов интерполяции x_i и коэффициентов c_i, а вычисление может быть сделано способом, аналогичным схеме Горнера.

"""
    InterpPolynomial <: AbstractPolynomial

Interpolation polynomials in Newton's form
"""
struct InterpPolynomial<:AbstractPolynomial
    degree::Int
    xval::NTuple{N, Float64} where N
    coeff::NTuple{N, Float64} where N
end

"""
    evpoly(p::Polynomial, z::Real)

Evaluate polynomial `p` at `z` using the Horner's rule
"""
function evpoly(p::InterpPolynomial, z::Real)
    ans = p.coeff[p.degree+1]
    for idx = p.degree:-1:1
        ans = ans * (z - p.xval[idx]) + p.coeff[idx]
    end
    return ans
end

Функция для вычисления значения полинома в обоих случаях называется одинаково — evpoly() — но принимает разные типы аргументов.

Кроме функции вычисления, неплохо бы ещё написать функцию, создающую полином по известным данным.

Для этого в Julia есть две методики: внешние конструкторы и внутренние конструкторы. Внешний конструктор — это просто функция, возвращающая объект соответствующего типа. Внутренний конструктор — это функция, которая вводится внутри описания структуры и заменяет собой стандартный конструктор. Для построения интерполяционных полиномов целесообразно использовать именно внутренний конструктор, поскольку

• получить полином удобнее не через узлы интерполяции и коэффициенты, а через узлы и значения интерполируемой функции
• узлы интерполяции должны обязательно быть различными
• число узлов и коэффициентов должно совпадать

Написание внутреннего конструктора, в котором гарантированно будут соблюдаться эти правила, гарантирует, что все создаваемые переменные типа InterpPolynomial, по крайней мере, могут корректно быть обработаны функцией evpoly().

Напишем конструктор обычных полиномов, принимающий на вход одномерный массив или кортеж коэффициентов. Конструктор интерполяционного полинома принимает на вход узлы интерполяции и желаемые значения в них и использует метод разделенных разностей для вычисления коэффициентов.

"""
    Polynomial <: AbstractPolynomial

Polynomials written in the canonical form

---

    Polynomial(v::T) where T<:Union{Vector{<:Real}, NTuple{<:Any, <:Real}})

Construct a `Polynomial` from the list of the coefficients. The coefficients are assumed to go from power 0 in the ascending order. If an empty collection is provided, the constructor returns a zero polynomial.
"""
struct Polynomial<:AbstractPolynomial
    degree::Int
    coeff::NTuple{N, Float64} where N
    function Polynomial(v::T where T<:Union{Vector{<:Real},
                                            NTuple{<:Any, <:Real}})
        # в случае пустого массива / кортежа в аргументе возвращаем P(x) ≡ 0
        coeff = isempty(v) ? (0.0,) : tuple([Float64(x) for x in v]...)
        # возврат значения - специальным оператором new
        # аргументы - перечисление значений полей
        return new(length(coeff)-1, coeff)
    end
end

"""
    InterpPolynomial <: AbstractPolynomial

Interpolation polynomials in Newton's form

---

    InterpPolynomial(xsample::Vector{<:Real}, fsample::Vector{<:Real})

Construct an `InterpPolynomial` from a vector of points `xsample` and corresponding function values `fsample`. All values in `xsample` must be distinct.
"""
struct InterpPolynomial<:AbstractPolynomial
    degree::Int
    xval::NTuple{N, Float64} where N
    coeff::NTuple{N, Float64} where N
    function InterpPolynomial(xsample::X,
                              fsample::F) where {X<:Union{Vector{<:Real},
                                                          NTuple{<:Any, <:Real}},
                                                 F<:Union{Vector{<:Real},
                                                          NTuple{<:Any, <:Real}}}
        # проверки на то, что все узлы различны, и значений f столько же, сколько узлов
        if !allunique(xsample)
            throw(DomainError("Cannot interpolate with duplicate X points"))
        end
        N = length(xsample)
        if length(fsample) != N
            throw(DomainError("Lengths of X and F are not the same"))
        end

        coeff = [Float64(f) for f in fsample]
        # алгоритм расчета разделенных разностей (Stoer, Bulirsch, Introduction to Numerical Analysis, гл. 2.1.3)
        for i = 2:N
            for j = 1:(i-1)
                coeff[i] = (coeff[j] - coeff[i]) / (xsample[j] - xsample[i])
            end
        end

        new(N-1, tuple([Float64(x) for x in xsample]...), tuple(coeff...))
    end
end

Кроме собственно генерации полиномов, неплохо бы иметь возможность производить с ними арифметические действия.

Поскольку в Julia арифметические операторы — это обычные функции, к которым в качестве синтаксического сахара добавлена инфиксная запись (выражения a + b и +(a, b) — оба допустимы и абсолютно идентичны), то перегрузка их делается точно так же, как и написание дополнительных методов к своим функциям.

Единственный тонкий момент — пользовательский код запускается из модуля (пространства имён) Main, а функции стандартной библиотеки находятся в модуле Base, поэтому при перегрузке нужно либо импортировать модуль Base, либо писать полное имя функции.

Итак, добавляем сложение полинома с числом:

# из-за особенностей парсера Base.+ не работает,
# и нужно писать Base.:+, что означает "символ :+ из модуля Base"
function Base.:+(p::Polynomial, x::Real)
    Polynomial(tuple(p.coeff[1] + x, p.coeff[2:end]...))
end

function Base.:+(p::InterpPolynomial, x::Real)
    # т.к. стандартный конструктор заменён на построение интерполяции по узлам и значениям - 
    # при сложении с числом нужно пересчитать значения во всех узлах.
    # Если операцию сложения планируется использовать часто -
    # стоит добавить конструктор по узлам и коэффициентам
    fval::Vector{Float64} = [evpoly(p, xval) + x for xval in p.xval]
    InterpPolynomial(p.xval, fval)
end

# чтобы сложение работало в любом порядке
function Base.:+(x::Real, p::AbstractPolynomial)
    return p + x
end

Для сложения двух обычных полиномов достаточно сложить коэффициенты, а при сложении интерполяционного полинома с другим можно найти значения суммы в нескольких точках и построить новую интерполяцию по ним.

function Base.:+(p1::Polynomial, p2::Polynomial)
    # при сложении нужно учесть, какой должна быть наивысшая степень
    deg = max(p1.degree, p2.degree)
    coeff = zeros(deg+1)

    coeff[1:p1.degree+1] .+= p1.coeff
    coeff[1:p2.degree+1] .+= p2.coeff

    Polynomial(coeff)
end

function Base.:+(p1::InterpPolynomial, p2::InterpPolynomial)
    xmax = max(p1.xval..., p2.xval...)
    xmin = min(p1.xval..., p2.xval...)
    deg = max(p1.degree, p2.degree)
    # для построения суммы строим чебышёвскую сетку между минимальным
    # и максимальным из узлов обоих полиномов
    xmid = 0.5 * xmax + 0.5 * xmin
    dx = 0.5 * (xmax - xmin) / cos(0.5 * π / (deg + 1))
    chebgrid = [xmid + dx * cos((k - 0.5) * π / (deg + 1)) for k = 1:deg+1]
    fsample = [evpoly(p1, x) + evpoly(p2, x) for x in chebgrid]
    InterpPolynomial(chebgrid, fsample)
end

function Base.:+(p1::InterpPolynomial, p2::Polynomial)
    xmax = max(p1.xval...)
    xmin = min(p1.xval...)
    deg = max(p1.degree, p2.degree)
    xmid = 0.5 * xmax + 0.5 * xmin
    dx = 0.5 * (xmax - xmin) / cos(0.5 * π / (deg + 1))
    chebgrid = [xmid + dx * cos((k - 0.5) * π / (deg + 1)) for k = 1:deg+1]
    fsample = [evpoly(p1, x) + evpoly(p2, x) for x in chebgrid]
    InterpPolynomial(chebgrid, fsample)
end

function Base.:+(p1::Polynomial, p2::InterpPolynomial)
    p2 + p1
end

Таким же образом можно добавить и другие арифметические операции над полиномами, в результате получив представление их в коде в естественной математической записи.

Пока на этом всё. Постараюсь дальше написать про реализацию других численных методов.

При подготовке использованы материалы:

1. Документация языка Julia: docs.julialang.org
2. Площадка обсуждения языка Julia: discourse.julialang.org
3. J.Stoer, W. Bulirsch. Introduction to Numerical Analysis
4. Хаб Julia: habr.com/ru/hub/julia
5. Think Julia: benlauwens.github.io/ThinkJulia.jl/latest/book.html