Define a new method with only a few changes

Question

I want to write a version that accepts a supplementary argument. The difference with the initial version only resides in a few lines of codes, potentially within loops. A typical example is to user a vector of weight w.

One solution is to completely rewrite a new function

function f(Vector::a)
   ...
   for x in a
      ...
      s += x[i]
      ... 
   end
   ...
end

function f(a::Vector, w::Vector)
   ... 
   for x in a
     ...
     s += x[i] * w[i]
     ...
   end
   ...
end

This solution duplicates code and therefore makes the program harder to maintain. I could split ... into different helper functions, which are called by both functions, but the resulting code would be hard to follow

Another solution is to write only one function and use a ? : structure for each line that should be changed

function f(a, w::Union(Nothing, Vector) = nothing)
   ....
   for x in a
      ...
      s += (w == nothing)? x[i] : x[i] * w[i]
      ...
   end
   ....
end

This code requires to check a condition at every step in a loop, which does not sound efficient, compared to the first version.

I'm sure there is a better solution, maybe using macros. What would be a good way to deal with this?

Answer 1

There are lots of ways to do this sort of thing, ranging from optional arguments to custom types to metaprogramming with @eval'ed code generation (this would splice in the changes for each new method as you loop over a list of possibilities).

I think in this case I'd use a combination of the approaches suggested by @ColinTBowers and @GnimucKey.

It's fairly simple to define a custom array type that is all ones:

immutable Ones{N} <: AbstractArray{Int,N}
    dims::NTuple{N, Int}
end
Base.size(O::Ones) = O.dims
Base.getindex(O::Ones, I::Int...) = (checkbounds(O, I...); 1)

I've chosen to use an Int as the element type since it tends to promote well. Now all you need is to be a bit more flexible in your argument list and you're good to go:

function f(a::Vector, w::AbstractVector=Ones(size(a))
    …

This should have a lower overhead than either of the other proposed solutions; getindex should inline nicely as a bounds check and the number 1, there's no type instability, and you don't need to rewrite your algorithm. If you're sure that all your accesses are in-bounds, you could even remove the bounds checking as an additional optimization. Or on a recent 0.4, you could define and use Base.unsafe_getindex(O::Ones, I::Int...) = 1 (that won't quite work on 0.3 since it's not guaranteed to be defined for all AbstractArrays).

Answer 2

In this case, using Optional Arguments may play the trick.

Just make the w argument default to ones().