CODEWITHSUNDEEP
functionmacrosclosuresjuliametaprogramming
Grayscale
Grayscale

Reputation: 1532

Julia create multiple slightly modified versions of a function

I have a function that looks like

function eom!(du, u, p)
    @views a, b = u[:,1], u[:,2];
    @views da, db = du[:,1], du[:,2];

    y = # some stuff involving p and a;
    da .= f(a, b, y);
    db .= g(b, a);
end

I now want to create a second a function that is the exact same, except the last line reads

db .= g(b, y);

How can I do this most cleanly? Of course, I could just copy and paste and give the functions slightly different names, but this seems unideal, especially if, as is plausible, I later want more functions where the second argument of g could be something else. Is there perhaps a way that I could pass into the function eom! an expression (via the argument p) that would specify the second argument of g? Or is there a way I could make some function eom_generator which can output all of the functions that I want? Perhaps macros are the central tool in doing this, but I am not sure.

Upvotes: 3

Views: 129

Answers (1)

phipsgabler
phipsgabler

Reputation: 20970

You could produce a closure:

function eom_generator(g)
    return function eom!(du, u, p)
        @views a, b = u[:,1], u[:,2]
        @views da, db = du[:,1], du[:,2]

        y = nothing # some stuff involving p and a;
        da .= f(a, b, y)
        db .= g(a, b, y)
    end
end

const eom1! = eom_generator((a, b, y) -> g(b, a))
const eom2! = eom_generator((a, b, y) -> g(b, y))

But since this is at the core of a differential equation, be sure to test whether you don't have any performance issues that way.

If you decide that you really need metaprogramming, you can use @eval in a loop:

for (i, expr) in enumerate((:(g(b, a)), :(g(b, y))))
    @eval function $(Symbol("eom", i, "!"))(du, u, p)
            @views a, b = u[:,1], u[:,2]
            @views da, db = du[:,1], du[:,2]

            y = nothing # some stuff involving p and a;
            da .= f(a, b, y)
            db .= $expr
        end
    end
end

Upvotes: 7

Related Questions