SNEWS2 / snewpy

A Python package for working with supernova neutrinos

Home Page:https://snewpy.readthedocs.io

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Vectorize probabilities in `FlavorTransformation`

Sheshuk opened this issue · comments

Currently FlavorTransformation.prob_ee and similar methods take time and energy as arguments.
These arguments can be arrays of values:

def prob_ex(self, t, E):
"""X -> e neutrino transition probability.
Parameters
----------
t : float or ndarray
List of times.
E : float or ndarray
List of energies.
Returns
-------
float or ndarray
Transition probability.

However the probabilities are often output as scalars, or as vectors only in E.

I suggest making the output universal: it should always be an array of shape (len(t), len(E)) (or scalar if both are scalars), even if the probability doesn't depend on E or t. This will make it easier for vectorized calculations, and more clear to the user.