Skip to content

Commit

Permalink
[#929] Temporary managing kph='ph' or kph='ph2' in _class8_plot_vos_s…
Browse files Browse the repository at this point in the history 
…pectro.py for cross- testing / validation, TBF
  • Loading branch information
@dvezinet
dvezinet committed May 28, 2024
1 parent 3a9b267 commit 92eac1f
Showing 1 changed file with 121 additions and 22 deletions.
143 changes: 121 additions & 22 deletions tofu/data/_class8_plot_vos_spectro.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@


# Built-in
import itertools as itt


# Common
Expand Down Expand Up @@ -354,12 +355,16 @@ def _prepare_ph(

# -----------------
# check up
# -----------------

kph = 'ph2'

if indlamb is None:
indlamb = int(dvos['lamb']['data'].size / 2)

# -----------------
# get mesh sampling
# -----------------

dsamp = coll.get_sample_mesh(
key=dvos['keym'],
Expand All @@ -380,6 +385,7 @@ def _prepare_ph(

# ------------------------------
# dhor
# ------------------------------

x0u = dsamp['x0']['data']
x1u = dsamp['x1']['data']
Expand Down Expand Up @@ -424,6 +430,10 @@ def _prepare_ph(

# ------------------------------
# prepare image in cross-section
# ------------------------------

# -----------------
# get shapes and initialize

n0, n1 = x0u.size, x1u.size
shape = (n0, n1)
Expand All @@ -440,12 +450,14 @@ def _prepare_ph(

lamb = np.full(shape, np.nan)
lambi = np.full(shape, np.nan)
lamb0 = np.full(shape, np.nan)
lambi0 = np.full(shape, np.nan)
dlamb = np.full(shape, np.nan)
dlambi = np.full(shape, np.nan)
cos = np.full(shape, np.nan)
cosi = np.full(shape, np.nan)

shape_cam = dvos['ph']['data'].shape[:-2]
shape_cam = dvos[kph]['data'].shape[:-2]
nc_cam = np.full(shape_cam, np.nan)
ph_cam = np.full(shape_cam, np.nan)
cos_cam = np.full(shape_cam, np.nan)
Expand All @@ -454,103 +466,163 @@ def _prepare_ph(

indr = dvos['indr_cross']['data']
indz = dvos['indz_cross']['data']
npts, nlamb = dvos['ph']['data'].shape[-2:]
ilambr = dvos['ilambr']['data']
npts = dvos[kph]['data'].shape[-2]

# -----------------
# multiply by dlamb
ph0 = dvos['ph']['data'] * np.mean(np.diff(dvos['lamb']['data']))

indch0 = indch
ph0 = dvos[kph]['data'] * np.mean(np.diff(dvos['lamb']['data']))
if is2d:
nc = dvos['ncounts']['data'].reshape((-1, npts))
ph = ph0.reshape((-1, npts, nlamb))
ph = ph0.reshape(tuple([-1] + list(ph0.shape[-2:])))
if kph == 'ph2':
ilambr = dvos['ilambr']['data'].reshape(tuple([-1] + list(ph0.shape[-2:])))
coss = dvos['cos']['data'].reshape((-1, npts))
lambc = dvos['lamb']['data'][None, None, None, :]
indch = indch[0] * shape_cam[1] + indch[1]

else:
nc = dvos['ncounts']['data']
ph = ph0
if kph == 'ph2':
ilambr = dvos['ilambr']['data']
coss = dvos['cos']['data']
lambc = dvos['lamb']['data'][None, None, :]

nci = nc[indch, :]
phi = ph[indch, :, :]
cossi = coss[indch, :]

# -----------------
# lambmax, lambmin
lambf = np.repeat(dvos['lamb']['data'][None, :], npts, axis=0)
for ni in shape_cam[::-1]:
lambf = np.repeat(lambf[None, :], ni, axis=0)

ind0 = ph0 == 0
lambf[ind0] = np.inf
lambmin = np.min(lambf, axis=-1)
lambf[ind0] = -np.inf
lambmax = np.max(lambf, axis=-1)
del lambf
shapel = ph0.shape[:-1]
lambmin = np.full(shapel, np.nan)
lambmax = np.full(shapel, np.nan)
if kph == 'ph2':
for ind in itt.product(*[range(ss) for ss in shapel]):
sli = tuple(list(ind) + [slice(None)])
iok = ph0[sli] > 0.
if np.any(iok):
if kph == 'ph2':
iok = dvos['ilambr']['data'][sli][iok]
lambmin[ind] = np.min(dvos['lamb']['data'][iok])
lambmax[ind] = np.max(dvos['lamb']['data'][iok])

lambmini = lambmin.reshape((-1, npts))[indch]
lambmaxi = lambmax.reshape((-1, npts))[indch]

# -----------------
# photon counts

nc_tot[indr, indz] = np.sum(nc, axis=0)
nc_toti[indr, indz] = np.sum(nci, axis=0)

# -----------------
# update indr, indz

iok = nc_tot[indr, indz] > 0.
iokn = np.nonzero(iok)[0]
ir, iz = indr[iok], indz[iok]

ioki = nc_toti[indr, indz] > 0.
iri, izi = indr[ioki], indz[ioki]

nc_tot[nc_tot == 0] = np.nan
nc_toti[nc_toti == 0] = np.nan

# -----------------
# dV

dV[ir, iz] = dvos['dV']['data'][iok]

# -----------------
# cos

cos[ir, iz] = np.nansum(nc[:, iok] * coss[:, iok], axis=0) / nc_tot[ir, iz]
cosi[iri, izi] = cossi[ioki]

# -----------------
# ph_tot

if np.any(np.isnan(ph)):
msg = "ph_count should not contain nans! (nansum copies)"
raise Exception(msg)

ph_tot[ir, iz] = np.sum(np.sum(ph[:, iok, :], axis=0), axis=-1)
ph_toti[iri, izi] = np.sum(phi[ioki, :], axis=-1)

# ------------------
# average wavelength
lamb[ir, iz] = np.sum(
np.sum(ph[:, iok, :] * dvos['lamb']['data'][None, None, :], axis=0),
axis=-1,
) / ph_tot[ir, iz]

for ii, (irii, izii) in enumerate(zip(ir, iz)):
if kph == 'ph2':
sli = (None, ilambr[:, iokn[ii], :])
else:
sli = (None, slice(None))
lamb[irii, izii] = (
np.sum(dvos['lamb']['data'][sli] * ph[:, iokn[ii], :])
/ ph_tot[irii, izii]
)

if kph == 'ph2':
ll = dvos['lamb']['data'][ilambr[indch, ioki, :]]
else:
ll = dvos['lamb']['data'][None, :]

lambi[iri, izi] = (
np.sum(phi[ioki, :] * dvos['lamb']['data'][None, :], axis=-1)
np.sum(phi[ioki, :] * ll, axis=-1)
/ ph_toti[iri, izi]
)

if kph == 'ph':
lamb0[ir, iz] = np.sum(
np.sum(ph[:, iok, :] * dvos['lamb']['data'][None, None, :], axis=0),
axis=-1,
) / ph_tot[ir, iz]

lambi0[iri, izi] = (
np.sum(phi[ioki, :] * dvos['lamb']['data'][None, :], axis=-1)
/ ph_toti[iri, izi]
)

assert np.allclose(lamb0, lamb, equal_nan=True), ('lamb\n', lamb, '\n', lamb0)
assert np.allclose(lambi0, lambi, equal_nan=True), ('lambi\n', lambi, '\n', lambi0)

# -----------------
# delta wavelength

dlamb[ir, iz] = (
np.nanmax(lambmax.reshape((-1, npts))[:, iok], axis=0)
- np.nanmin(lambmin.reshape((-1, npts))[:, iok], axis=0)
)
dlambi[iri, izi] = lambmaxi[ioki] - lambmini[ioki]

# -----------------
# adjust

ph_tot[ph_tot == 0.] = np.nan
ph_toti[ph_toti == 0.] = np.nan

# -----------------------
# prepare image on camera
# -----------------------

nc_cam[...] = np.nansum(dvos['ncounts']['data'], axis=-1)
iok = nc_cam > 0.
nc_cam[~iok] = np.nan

ph_cam[iok] = np.nansum(np.nansum(ph0, axis=-1), axis=-1)[iok]

if kph == 'ph2':
lamb_aa = dvos['lamb']['data'][dvos['ilambr']['data']]
else:
lamb_aa = lambc

lamb_cam[iok] = (
np.nansum(np.nansum(ph0 * lambc, axis=-1), axis=-1)[iok]
np.nansum(np.nansum(ph0 * lamb_aa, axis=-1), axis=-1)[iok]
/ ph_cam[iok]
)

Expand All @@ -567,20 +639,47 @@ def _prepare_ph(

# ----------------------
# prepare per wavelength
# ----------------------

if kph == 'ph2':
indlamb2 = dvos['ilambr']['data'] == indlamb
ii2 = np.any(indlamb2, axis=-2)
print(ii2.shape, ii2.sum(), ph.shape)
print()
print(np.sum(np.any(ii2, axis=-1)), np.sum(~np.any(ii2, axis=-1)))
ii2[~np.any(ii2, axis=-1), 0] = True
print(ii2.sum())
assert np.all(np.sum(indlamb2, axis=-2) <= 1)
assert ii2.sum() == ph.shape[0], (ii2.shape, ii2.size, ii2.sum())
print(indlamb2.shape, indlamb2.sum(), ii2.shape)
print(ph0.shape, ph.shape, phi.shape)
ph_cam_lamb = np.nansum(ph0, axis=-1, where=indlamb2)
ph_cam_lamb2 = np.nansum(ph0, axis=-2, where=indlamb2)[ii2]
print(ph_cam_lamb.shape)
print(ph_cam_lamb2.shape)
indlamb2 = ilambr == indlamb
assert np.all(np.sum(indlamb2, axis=-1) <= 1)
print(indlamb2.shape)
ph_tot_lamb[indr, indz] = np.nansum(ph, axis=0, where=indlamb2)
ph_toti_lamb[indr, indz] = np.nansum(phi, axis=0, where=indlamb2)
print(ph_cam_lamb.shape, ph_tot_lamb.shape, ph_toti_lamb.shape)

ph_cam_lamb = np.nansum(ph0[..., indlamb], axis=-1)
ph_tot_lamb[indr, indz] = np.nansum(ph[..., indlamb], axis=0)
ph_toti_lamb[indr, indz] = np.nansum(phi[..., indlamb], axis=0)
else:
ph_cam_lamb = np.nansum(ph0[..., indlamb], axis=-1)
ph_tot_lamb[indr, indz] = np.nansum(ph[..., indlamb], axis=0)
ph_toti_lamb[indr, indz] = np.nansum(phi[..., indlamb], axis=0)

ph_cam_lamb[ph_cam_lamb == 0] = np.nan

# ----------------------
# delta_lamb for etendue
# ----------------------

delta_lamb = coll.get_diagnostic_data(key, data='dlamb')[0][key_cam[0]]

# -------------------
# extent
# -------------------

x0 = dsamp['x0']['data']
dx0 = x0[1] - x0[0]
Expand Down

0 comments on commit 92eac1f

Please sign in to comment.