Complex number class in Lua

Question

This is my first attemp with Lua.

I decided to create this class, because it have lots of "tweaks" and "tricks".

I want to know if I did it as best practices Lua ways.

My Lua version is 5.4.1

require "math"

Complex = {
    IMAGINARY_CHAR = "i"
}

function Complex:new(r, i)
    o = {
        r = r,
        i = i
    }

    setmetatable(o, self)

    self.__index = self

    return o
end

function Complex:__newindex(r, i)
end

function Complex:__unm()
    return Complex:new(
        -self.r,
        -self.i
    )
end

function Complex:__add(other)
    if type(other) == "number" then
        return Complex:new(
            self.r + other,
            self.i
        )
    else
        return Complex:new(
            self.r + other.r,
            self.i + other.i
        )
    end
end

function Complex:__sub(other)
    return self:__add(- other);
end

function Complex:__mul(other)
    if type(other) == "number" then
        return Complex:new(
            self.r * other,
            self.i * other
        )
    else
        return Complex:new(
            self.r * other.r - self.i * other.i,
            self.i * other.r + self.r * other.i
        )
    end
end

function Complex:__eq(other)
    return
        self.r == other.r and
        self.i == other.i
end

function Complex:__lt(other)
    -- incorrect but for sorting
    if self.r == other.r then
        return self.i < other.i
    else
        return self.r < other.r
    end
end

function Complex:tostring(i)
    if type(i) ~= "string" then
        i = "i"
    end

    if self.r == 0 and self.i == 0 then
        return "( 0 )"
    elseif self.i == 0 then
        return "( " .. self.r .. " )"
--  elseif self.r == 0 then
--      return "( " .. i .. self.i .. " )"
    else
        return "( " .. self.r .. " + " .. i .. self.i .. " )"
    end
end

function Complex:__tostring(i)
    return self:tostring(self.IMAGINARY_CHAR)
end

function Complex:abs2()
    return self.r * self.r + self.i * self.i
end

function Complex:abs()
    return math.sqrt(self:abs2())
end

-- ===================================================================

Complex.IMAGINARY_CHAR = "j" -- Electrical notation

x = Complex:new(0, 10)
y = Complex:new(-5, 5)

print(x)
print(y)

print(x:abs())
print(y:abs())


print(x + y)
print(x - y)

print(x + 12)
print(x - 12)

print(x == y)
print(x ~= y)

print(x ~= x)
print(x ~= x)

print(x < y)
print(x > y)

print(x <= y)
print(x >= y)

Answer 1

Possible improvements:

• localise functions of the math library, and sometimes, self.r and self.i. It will reduce the number of table lookups, improving the performance, and will make the expressions simpler,
• define __call metamethod for Complex. It will allow you to replace Complex:new(x, y) with Complex(x, y),
• the 'constructor' should correctly handle cases when it is called with one parameter (Complex(5) = 5 + 0i). It can be achieved with a 'nullsafe operator' emulated as local a = user_input or value_if_null,
• in can also be mage more concise, using the fact that setmetatable returns the table,
• __newindex metamethod is not needed here, unless you want to set the absolute value or phase with code like c.abs = 1 or code.arg = pi. In addition, the arguments to __newindex are the table, the absent key and value set to it,
• you can emulate the ternary operator in Lua with and and or, which allows to write more concise code: local a = condition and on_success or on_failure,
• your arithmetic metamethods should handle the cases when the first argument (self) is a number,
• you can define conjugate method and __div metamethod for division,
• to square a table element, I'd rather recommend using ^ rather than multiplying by self. It will save table lookups,
• you can define methods for polar coordinates and __pow metamethod for Euler's and de Moivre's formulae,
• the __lt metamethod can be simplified using boolean operations,
• Complex:tostring() can be simplified using 'ternary' operators, also handling of negative imaginary parts can be improved,
• whether to put the i before or after the imaginary part and whether to surround a complex number with parentheses can be customisable,
• the expression tests can be automated somewhat with load() function.

The improved code:

complex.lua:

local math = require 'math'
local pi, sin, cos, atan2 = math.pi, math.sin, math.cos, math.atan2
local sqrt, exp, ln = math.sqrt, math.exp, math.log
local floor, abs = math.floor, math.abs

complex = {
    IMAGINARY_CHAR = 'i'
}

function complex:new (r, i)
    self.__index = self
    return setmetatable (type (r) == 'table' and r or {r = r, i = i or 0}, self)
end

setmetatable (complex, {
    __call = function (tbl, r, i)
        return complex:new (r, i)
    end}
)

function complex:__unm ()
    return complex (-self.r, -self.i)
end

function complex:__add (other)
    return type (self) == 'number'   and complex (self + other.r, other.i)
        or type (other) == 'number'  and complex (self.r + other, self.i)
                                     or complex (self.r + other.r, self.i + other.i)
end

function complex:__sub (other)
    return self + (-other);
end

function complex:__mul (other)
    return type (self) == 'number'   and complex (self * other.r, self * other.i)
        or type (other) == 'number'  and complex (self.r * other, self.i * other)
                                     or complex (self.r * other.r - self.i * other.i, self.i * other.r + self.r * other.i)
end

function complex:conjugate ()
    return complex (self.r, -self.i)
end

function complex:__div (denominator)
    -- https://www.mesacc.edu/~scotz47781/mat120/notes/complex/dividing/dividing_complex.html
    local conjugate = type (denominator) == 'number' and complex (denominator) or denominator:conjugate ()
    local new_numerator, new_denominator = self * conjugate, denominator * conjugate
    -- new_denominator is real.
    return complex (new_numerator.r / new_denominator.r, new_numerator.i / new_denominator.r)
end

function complex:abs2 ()
    return self.r ^ 2 + self.i ^ 2
end

function complex:abs ()
    return sqrt (self:abs2 ())
end

function complex:polar (abs, arg)
    return complex (abs * cos (arg), abs * sin (arg))
end

function complex:arg ()
    return atan2 (self.i, self.r)
end

function complex:exp ()
    local abs = exp (self.r)
    return complex (abs * cos (self.i), abs * sin (self.i))
end
    
function complex:__pow (power)
    -- Euler:
    local x = type (self) == 'number' and self or self.i == 0 and self.r > 0 and self.r or nil
    if x then
        return complex (power * ln (x)):exp ()
    else
    -- de Moivre:
        local n = type (power) == 'number' and power or power.i == 0 and power.r or nil
        if n and floor (n) == n then
            local abs, arg = self:abs (), self:arg ()
            return complex:polar (abs ^ n, n * arg)
        end
    end
end

function complex:__eq (other)
    return
        self.r == other.r and
        self.i == other.i
end

function complex:__lt (other)
    -- incorrect but for sorting:
    return self.r < other.r or self.r == other.r and self.i < other.i
end

function complex:tostring (i, prefix, parentheses)
    local im = type (i) == 'string' and i or 'i'
    local r, i = self.r, self.i

    local r_str = (r ~= 0 or i == 0) and r or ''
    local i_str = i ~= 0 and (prefix and im .. abs (i) or (abs (i) == 1 and '' or abs (i)) .. im) or ''
    local space = (r ~= 0 and i ~= 0) and ' ' or ''
    local sign = r ~= 0 and i > 0 and '+' or i < 0 and '-' or ''
    
    local str = r_str .. space .. sign .. space .. i_str
    if parentheses then
        str = '(' .. str .. ')'
    end
    return str
end

function complex:__tostring ()
    return self:tostring (self.IMAGINARY_CHAR, self.PREFIX, self.PARENTHESES)
end

return complex

test.lua:

local load = _VERSION == 'Lua 5.1' and function (chunk, _, __, context)
    for key, value in pairs (context) do
        _G [key] = value
    end
    return loadstring (chunk)
end or load

local complex = require 'complex'

complex.IMAGINARY_CHAR = 'j' -- electrical notation.
complex.PREFIX = true
complex.PARENTHESES = true

local context = {
    complex = complex,
    x = complex (0, 10),
    y = complex (-5, 5),
    z = complex (5),
    pi = math.pi
}

local cases = {
    'x', 'y', 'z', 'x:abs ()',
    'x + y', 'x - y', 'x + 12', 'x - 12', '12 + x', '12 - x',
    'x * y', 'y * 12', '2 * x',
    'y:conjugate ()', 'x / y', 'y / y', 'x / 5', '5 / x', '5 / x * x',
    'y:abs ()', 'y:arg () / pi * 180', 'complex:polar (2, pi / 4)', 'y:exp ()', 'x ^ 2', '2 ^ y', '2 ^ x',
    'x == y', 'x ~= y', 'x < y', 'x > y', 'x <= y', 'x >= y'
}

for _, expr in ipairs (cases) do
    print (expr, assert (load ('return ' .. expr, nil, 't', context))())
end

For further improvements, see my repository.