new version of snow_impurities slower

[BaptisteVandecrux]

    t1 = time.process_time()
    ntype, bf, conc = snow_impurities_new(snow.alb_sph, snow.bal)
    t2 = time.process_time()
    print(t2-t1)
    t1 = time.process_time()
    ntype_old, bf_old, conc_old =  snow_impurities(snow.alb_sph.unstack(dim='xy').values, snow.bal.unstack(dim='xy').values)
    t2 = time.process_time()
    print(t2-t1)

5.0625
0.65625

def snow_impurities(alb_sph, bal):
        # analysis of snow impurities
    # ( the concentrations below 0.0001 are not reliable )        
    # bf    normalized absorption coefficient of pollutants ay 1000nm ( in inverse mm)
    # bm    Angstroem absorption coefficient of pollutants ( around 1 - for soot, 3-7 for dust)
    bm=np.empty(bal.shape)*np.nan
    bf=np.empty(bal.shape)*np.nan
    p1 = np.empty(bal.shape)*np.nan       
    p2 = np.empty(bal.shape)*np.nan

    ind_nonan = np.logical_and(np.logical_not(np.isnan(alb_sph[0,:,:])),
                               np.logical_not(np.isnan(alb_sph[1,:,:])))   
    p1[ind_nonan]=np.log(alb_sph[0,ind_nonan])**2
    p2[ind_nonan]=np.log(alb_sph[1,ind_nonan])**2
    bm[ind_nonan]=np.log(p1[ind_nonan]/p2[ind_nonan])/np.log(w[1]/w[0])

    # type of pollutants
    ntype=np.nan*bal
    ntype[bm <= 1.2]=1   # soot
    ntype[bm > 1.2]=2    # dust

    soda = bm*np.nan
    soda[bm>=0.1]=(w[0])**bm[bm>=0.1]
    bf=soda*p1/bal

    # normalized absorption coefficient of pollutants at the wavelength  1000nm
    k_abs_0=p1/bal
    # bal   -effective absorption length in microns

    B_soot=1.6        # enhancement factors for soot
    B_ice= 0.9       # enhancement factors for ice grains
    alfa_soot=4.*np.pi*0.47/w[0]  # bulk soot absorption coefficient at 1000nm
    k_dust=0.01       # volumetric absorption coefficient of dust

    conc = bal*np.nan
    conc[ntype == 1] = B_soot*k_abs_0[ntype == 1]/B_ice/alfa_soot
    conc[ntype == 2] = B_soot*k_abs_0[ntype == 2]/k_dust
    ntype[bm <= 0.5] = 3 # type is other or mixture
    ntype[bm >= 10.] = 4 # type is other or mixture
    return ntype, bf, conc

def snow_impurities_new(alb_sph, bal):
    # analysis of snow impurities
    # ( the concentrations below 0.0001 are not reliable )
    # bf    normalized absorption coefficient of pollutants ay 1000nm ( in inverse mm)
    # bm    Angstroem absorption coefficient of pollutants ( around 1 - for soot, 3-7 for dust)

    ind_nonan = ~np.isnan(alb_sph.sel(band=0)) & ~np.isnan(alb_sph.sel(band=1))

    p1 = np.log(alb_sph.sel(band=0))**2
    p2 = np.log(alb_sph.sel(band=1))**2
    bm = np.log(p1 / p2) / np.log(wls.sel(band=1) / wls.sel(band=0))

    # type of pollutants
    ntype = 1 + (bm > 1.2)     # 1=soot, 2 = dust

    soda = (wls.sel(band=0)**bm).where(bm >= 0.1)
    bf = soda * p1 / bal

    # normalized absorption coefficient of pollutants at the wavelength  1000nm
    k_abs_0 = p1 / bal
    # bal   -effective absorption length in microns

    B_soot = 1.6        # enhancement factors for soot
    B_ice = 0.9       # enhancement factors for ice grains
    alfa_soot = 4. * np.pi * 0.47 / wls.sel(band=0)  # bulk soot absorption coefficient at 1000nm
    k_dust = 0.01       # volumetric absorption coefficient of dust

    # conc[ntype == 1] = B_soot*k_abs_0[ntype == 1] / B_ice / alfa_soot
    # conc[ntype == 2] = B_soot*k_abs_0[ntype == 2] / k_dust

    conc = (ntype == 1) * B_soot * k_abs_0 / B_ice / alfa_soot + \
           (ntype == 2) * B_soot * k_abs_0 / k_dust

    ntype = ntype.where(bm > 0.5, 3)  # type is other or mixture
    ntype = ntype.where(bm < 10., 4)  # type is other or mixture

    return ntype.where(ind_nonan), bf.where(ind_nonan), conc.where(ind_nonan)