SNEWS2 / snewpy

A Python package for working with supernova neutrinos
https://snewpy.readthedocs.io
BSD 3-Clause "New" or "Revised" License
26 stars 19 forks source link

Allow scalar input for 't', 'E' and 'data' arguments of Container class #301

Closed Sheshuk closed 7 months ago

Sheshuk commented 7 months ago

Fixing problem if a user passes a scalar value to container. Example:

import numpy as np
import astropy.units as u
from snewpy.flux import Flux
from snewpy.neutrino import Flavor

# this was the only valid way before this fix: wrap everything in []
Flux(data = [1] << u.Unit('1/(MeV s cm**2)'), 
     flavor=[Flavor.NU_E],
     time=[1]<<u.s,
     energy=[1]<<u.MeV)

#now it's possible to pass scalars if there is only a single value or point
Flux(data = 1 << u.Unit('1/(MeV s cm**2)'), 
     flavor=Flavor.NU_E, 
     time=1<<u.s,
     energy=1<<u.MeV)

#they're converted internally to arrays. This produces 
# d2FdEdT (1, 1, 1) [1 / (MeV s m2)]: <1 flavor(0;0) x 1 time(1.0 s;1.0 s) x 1 energy(1.0 MeV;1.0 MeV)>
JostMigenda commented 7 months ago

Even worse—before, you needed data = [[[1]]] << u.Unit('1/(MeV s cm**2)') (with triple brackets). I agree this is a nice improvement; thanks!

Sheshuk commented 7 months ago

Thanks, you're right! I didn't think about that. That's definitely against what the user would expect.

Sheshuk commented 7 months ago

I added:

  1. Reshaping data array if it's 1D to one of the expected shapes.
  2. Array shape validation vs. expected shapes.

Maybe it still doesn't cover all the possible corner cases. But in general this class is not supposed to be created by user, but rather produced in the SupernovaModel.get_flux etc...

But I need to think of more strict conditions, when to use this reshaping