What is the best way to define a function that needs to be passed to another function?

Question

Please consider the following code:

def build_matrix(x_coordinates, y_coordinates, calc_element, **kwargs):

    matrix = np.zeros((len(x_coordinates), len(y_coordinates)))
    for i in range(len(x_coordinates)):
        for j in range(len(y_coordinates)):
            matrix += calc_element(x_coordinates[i], y_coordinates[j], i, j, **kwargs)

    return matrix

The function build_matrix is used to build a matrix by filling up its elements using a function called calc_element. Since I don't know exactly what all the possible functions are that might be passed to this build_matrix function, I use optional keyword arguments so that I don't need to change build_matrix every time that I want to use it with another function (the only thing that all these functions have in common is that they need the arguments x_coords_i, y_coords_j, i, j)

However, I was wondering whether it would be Pythonic to use a class in this case instead of a function. For example, if I would define a particular version of calc_element in the following way

class calc_element_velocity_mass():
    def __init__(self, velocity, mass):
        self.velocity = velocity
        self.mass = mass

    def __call__(self, x_coords_i, y_coords_j, i, j):
        return self.velocity * i / x_coords_i - y_coords_j * j * self.mass

then after initialising it as follows

calc_element = calc_element_velocity_mass(20, 10)

I would be able to pass it to build_matrix without using optional keyword arguments, i.e., without passing for example kwargs = dict(mass = 10, velocity = 20) to build_matrix.

Question: Would this be a proper way of classes? And is my original implementation correct, or should I approach this in a different way? Thanks!

Answer 1

Your class solution would definitely work, but a closure would be a much easier way to do the exact same thing:

def calc_element_velocity_mass(velocity, mass):
    def f(x_coords_i, y_coords_j, i, j):
        return velocity * i / x_coords_i - y_coords_j * j * mass
    return f

You could also return a lambda instead of f, but I used a named function here for clarity.