Should find a pattern that allows more flexible random distributions

[KevinMcCarthyAtIDM]

Right now, we hard-code specific random distributions - e.g., "inf_mean", "inf_std" for a normally distributed infectious distribution. In the spirit of "composing" a model, I suppose this isn't the worst, as a user should be able to make this choice and move on.
However, I think it would be an improvement to find a pattern that allowed more flexibility - being able to pass distributions around as function handles or something. I think it's tricky to do that with Numba, maybe with GPUs as well, or anywhere where we lose compile-time optimization by not knowing the run-time signature of the function, but long-run if we can find a pattern that supports both, that would be great.

[clorton]

Agreed. I need to look into the possibility of passing function pointers into Numba compiled methods.

[clorton]

This works, so we should be able to pass distributions into the Numba-fied methods that need flexibility:

import numba as nb
import numpy as np

@nb.njit((nb.int32,), cache=True)
def add_answer(x):
    return x + 42

func_type = nb.types.FunctionType(nb.int32(nb.int32))

@nb.njit((nb.int32[:], func_type), parallel=True, cache=True)
def apply_func(arr, func):
    result = np.empty_like(arr)
    for i in nb.prange(arr.size):
        result[i] = func(arr[i])
    return result

# Test the composed functions
data = np.array([1, 2, 3, 4], dtype=np.int32)
output = apply_func(data, add_answer)
print(output)  # Output: [2 3 4 5]

[KevinMcCarthyAtIDM]

Thanks, Chris! This pattern looks like it would work. The pattern from the user side that I'm imagining, then, would look something like

laser/src/distributions.py

import numba as nb
import numpy as np
@nb.njit((nb.float32,), cache=True)
def nbexponential(params):
    return np.random.exponential(params[0])

in model setup/execution code

from laser.distributions import nbexponential
parameters = PropertySet({"infectious_period": {"distribution": nbexponential, "params": [scale]})

and in transmission.py

def nb_transmission_update(..., distribution, params):  
for i in nb.prange(count):
     ...
                    itimers[i] =distribution(params)
     ...
        return

nb_transmission_update(..., params.infectious_period.distribution, params.infectious_period.params)

Want to be flexible with respect to number of input parameters.
I don't know how much we would need flexibility w/r/t function signature. I imagine that most places we are drawing randoms, we can know ahead what the signature will be, and there aren't many to cover. :

float(floats).#Continuous distributions with continuous parameters: Exponential, gaussian, lognormal, ...etc.
int(floats). #Discrete distributions with continuous parameters only.  E.g., Poisson
int(ints) #Discrete distributions with discrete parameters - hypergeometric-family distributions
int(ints, floats). #Discrete distributions with integer parameters and continuous parameters: binomial family, e.g.

Anyway, this was just the sort of pattern I came up with thinking about it for a short time, that would let a user easily specify the input distribution as a parameter. Probably still some stuff to be worked out, but interested to try this out soon.

[clorton]

I think you are on the right track. I think the first bit of code, nbexponential(), needs a minor change - the parameter signature should be nb.float32[:] (note the array syntax).

Also, I would imagine that float32 would cover all our parameter needs - floating point and integer.

You may be correct that we need to differentiate between the return value types - floating point or integer, although we might try getting away with floating point there as well and letting the calling function cast to integer as appropriate - but that might hurt my OCD knowing that a function returning integers was getting them converted to float and then back to integer! (TBD)