PythonicDISORT 1.1__py3-none-any.whl → 1.3__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (11) hide show
  1. PythonicDISORT/_assemble_intensity_and_fluxes.py +151 -130
  2. PythonicDISORT/_solve_for_coeffs.py +208 -142
  3. PythonicDISORT/_solve_for_gen_and_part_sols.py +84 -94
  4. PythonicDISORT/pydisort.py +200 -184
  5. PythonicDISORT/subroutines.py +191 -137
  6. {pythonicdisort-1.1.dist-info → pythonicdisort-1.3.dist-info}/METADATA +3 -3
  7. pythonicdisort-1.3.dist-info/RECORD +11 -0
  8. {pythonicdisort-1.1.dist-info → pythonicdisort-1.3.dist-info}/WHEEL +1 -1
  9. pythonicdisort-1.1.dist-info/RECORD +0 -11
  10. {pythonicdisort-1.1.dist-info → pythonicdisort-1.3.dist-info}/LICENSE.md +0 -0
  11. {pythonicdisort-1.1.dist-info → pythonicdisort-1.3.dist-info}/top_level.txt +0 -0
PythonicDISORT/_assemble_intensity_and_fluxes.py CHANGED
@@ -9,14 +9,15 @@ def _assemble_intensity_and_fluxes(
9
9
  scaled_omega_arr, # Delta-scaled single-scattering albedos
10
10
  tau_arr, # Lower boundary of layers
11
11
  scaled_tau_arr_with_0, # Delta-scaled lower boundary of layers with 0 inserted in front
12
- mu_arr_pos, mu_arr, # Quadrature nodes for 1) upper 2) both hemispheres
13
- M_inv, W, # 1) 1 / mu; 2) quadrature weights for each hemisphere
12
+ mu_arr_pos, # Upper hemisphere quadrature nodes (cosine polar angles)
13
+ M_inv, W, # 1) 1 / `mu_arr_pos`; 2) quadrature weights for each hemisphere
14
14
  N, NQuad, NLeg, # Number of 1) upper 2) both hemispheres quadrature nodes; 3) phase function Legendre coefficients
15
15
  NFourier, NLayers, NBDRF, # Number of 1) intensity Fourier modes; 2) layers; 3) BDRF Fourier modes
16
16
  is_atmos_multilayered, # Is the atmosphere multilayered?
17
17
  weighted_scaled_Leg_coeffs, # Weighted and delta-scaled Legendre coefficients
18
18
  BDRF_Fourier_modes, # BDRF Fourier modes
19
- mu0, I0, rescale_factor, phi0, # Properties of the direct beam
19
+ mu0, I0, I0_div_4pi, # Primary direct beam parameters
20
+ rescale_factor, phi0, # Secondary direct beam parameters
20
21
  there_is_beam_source, # Is there a beam source?
21
22
  b_pos, b_neg, # Dirichlet BCs
22
23
  b_pos_is_scalar, b_neg_is_scalar, # Is each Dirichlet BCs scalar and so isotropic?
@@ -47,7 +48,6 @@ def _assemble_intensity_and_fluxes(
47
48
  | `tau_arr` | `NLayers` |
48
49
  | `scaled_tau_arr_with_0` | `NLayers + 1` |
49
50
  | `mu_arr_pos` | `NQuad/2` |
50
- | `mu_arr` | `NQuad` |
51
51
  | `M_inv` | `NQuad/2` |
52
52
  | `W` | `NQuad/2` |
53
53
  | `N` | scalar |
@@ -61,6 +61,7 @@ def _assemble_intensity_and_fluxes(
61
61
  | `BDRF_Fourier_modes` | `NBDRF` |
62
62
  | `mu0` | scalar |
63
63
  | `I0` | scalar |
64
+ | `I0_div_4pi` | scalar |
64
65
  | `rescale_factor` | scalar |
65
66
  | `phi0` | scalar |
66
67
  | `there_is_beam_source` | scalar |
@@ -83,7 +84,7 @@ def _assemble_intensity_and_fluxes(
83
84
  | ----------------- | -------------------------------------- |
84
85
  | `G_collect` | `NFourier x NLayers x NQuad x NQuad` |
85
86
  | `G_collect_0` | `NLayers x NQuad x NQuad` |
86
- | `G_inv_collect_0` | `NLayers x NQuad x NQuad` or `None` |
87
+ | `G_inv_mu_inv` | `NLayers x NQuad x NQuad` or `None` |
87
88
  | `K_collect` | `NFourier x NLayers x NQuad` |
88
89
  | `K_collect_0` | `NLayers x NQuad` |
89
90
  | `B_collect` | `NFourier x NLayers x NQuad` or `None` |
@@ -108,45 +109,45 @@ def _assemble_intensity_and_fluxes(
108
109
  outputs = _solve_for_gen_and_part_sols(
109
110
  NFourier,
110
111
  scaled_omega_arr,
111
- mu_arr_pos, mu_arr,
112
+ mu_arr_pos,
112
113
  M_inv, W,
113
114
  N, NQuad, NLeg,
114
115
  NLayers,
115
116
  weighted_scaled_Leg_coeffs,
116
- mu0, I0,
117
+ mu0, I0_div_4pi,
117
118
  there_is_beam_source,
118
- Nscoeffs,
119
119
  there_is_iso_source,
120
120
  )
121
121
  if there_is_beam_source and there_is_iso_source:
122
122
  G_collect, K_collect, B_collect, G_inv_collect_0 = outputs
123
123
  B_collect_0 = B_collect[0, :, :]
124
+ G_inv_mu_inv = G_inv_collect_0 @ np.concatenate([M_inv, -M_inv])
124
125
  elif there_is_beam_source and not there_is_iso_source:
125
126
  G_collect, K_collect, B_collect = outputs
126
127
  B_collect_0 = B_collect[0, :, :]
127
- G_inv_collect_0 = None
128
+ G_inv_mu_inv = None
128
129
  elif not there_is_beam_source and there_is_iso_source:
129
130
  G_collect, K_collect, G_inv_collect_0 = outputs
130
131
  B_collect = None
132
+ G_inv_mu_inv = G_inv_collect_0 @ np.concatenate([M_inv, -M_inv])
131
133
  else:
132
134
  G_collect, K_collect = outputs
133
135
  B_collect = None
134
- G_inv_collect_0 = None
135
-
136
+ G_inv_mu_inv = None
136
137
 
137
138
  GC_collect = _solve_for_coeffs(
138
139
  NFourier,
139
140
  G_collect,
140
141
  K_collect,
141
142
  B_collect,
142
- G_inv_collect_0,
143
+ G_inv_mu_inv,
143
144
  scaled_tau_arr_with_0,
144
- mu_arr, mu_arr_pos, mu_arr_pos * W,
145
+ mu_arr_pos, mu_arr_pos * W,
145
146
  N, NQuad,
146
147
  NLayers, NBDRF,
147
148
  is_atmos_multilayered,
148
149
  BDRF_Fourier_modes,
149
- mu0, I0,
150
+ mu0, I0_div_4pi,
150
151
  there_is_beam_source,
151
152
  b_pos, b_neg,
152
153
  b_pos_is_scalar, b_neg_is_scalar,
@@ -166,7 +167,7 @@ def _assemble_intensity_and_fluxes(
166
167
 
167
168
  # Construct the intensity function (refer to section 3.7 of the Comprehensive Documentation)
168
169
  # --------------------------------------------------------------------------------------------------------------------------
169
- def u(tau, phi, is_antiderivative_wrt_tau=False, return_Fourier_error=False, return_tau_arr=False):
170
+ def u(tau, phi, is_antiderivative_wrt_tau=False, return_Fourier_error=False, return_tau_arr=False, *, _return_l=False):
170
171
  """
171
172
  Intensity function with arguments `(tau, phi)` of types `(array or float, array or float)`.
172
173
  Returns an ndarray with axes corresponding to variation with `mu, tau, phi` respectively.
@@ -182,7 +183,7 @@ def _assemble_intensity_and_fluxes(
182
183
 
183
184
  # Atmospheric layer indices
184
185
  l = np.argmax(tau[:, None] <= tau_arr[None, :], axis=1)
185
- scaled_tau_arr_l = scaled_tau_arr_with_0[l + 1]
186
+ scaled_tau_arr_l = scaled_tau_arr_with_0[1:][l]
186
187
  scaled_tau_arr_lm1 = scaled_tau_arr_with_0[l]
187
188
 
188
189
  # Delta-M scaling
@@ -201,60 +202,65 @@ def _assemble_intensity_and_fluxes(
201
202
  axis=-1,
202
203
  )
203
204
 
204
- if is_antiderivative_wrt_tau and there_is_beam_source:
205
+ direct_beam_contribution = 0
206
+ if is_antiderivative_wrt_tau:
207
+ if there_is_beam_source:
208
+ direct_beam_contribution = (
209
+ B_collect.transpose(0, 2, 1) / (-scale_tau / mu0)[None, None, :]
210
+ )[:, :, l] * np.exp(-scaled_tau / mu0)[None, None, :]
211
+
205
212
  um = (
206
213
  np.einsum(
207
- "mtij, mtj -> mti",
214
+ "mtij, mtj -> mit",
208
215
  GC_collect[:, l, :, :],
209
216
  np.exp(exponent) / (scale_tau[None, :, None] * K_collect)[:, l, :],
210
217
  optimize=True,
211
218
  )
212
- + (B_collect / (-scale_tau / mu0)[None, :, None])[:, l, :]
213
- * np.exp(-scaled_tau / mu0)[None, :, None]
219
+ + direct_beam_contribution
214
220
  )
215
- elif is_antiderivative_wrt_tau and not there_is_beam_source:
216
- um = np.einsum(
217
- "mtij, mtj -> mti",
218
- GC_collect[:, l, :, :],
219
- np.exp(exponent) / (scale_tau[None, l, None] * K_collect)[:, l, :],
220
- optimize=True,
221
- )
222
- elif not is_antiderivative_wrt_tau and there_is_beam_source:
223
- um = np.einsum(
224
- "mtij, mtj -> mti", GC_collect[:, l, :, :], np.exp(exponent), optimize=True
225
- ) + B_collect[:, l, :] * np.exp(-scaled_tau / mu0)[None, :, None]
226
221
  else:
227
- um = np.einsum(
228
- "mtij, mtj -> mti", GC_collect[:, l, :, :], np.exp(exponent), optimize=True
222
+ if there_is_beam_source:
223
+ direct_beam_contribution = (
224
+ B_collect.transpose(0, 2, 1)[:, :, l]
225
+ * np.exp(-scaled_tau / mu0)[None, None, :]
226
+ )
227
+
228
+ um = (
229
+ np.einsum(
230
+ "mtij, mtj -> mit", GC_collect[:, l, :, :], np.exp(exponent), optimize=True
231
+ )
232
+ + direct_beam_contribution
229
233
  )
230
234
 
231
235
  # Contribution from particular solution for isotropic internal sources
232
236
  if there_is_iso_source:
233
237
  l_uniq, l_inv = np.unique(l, return_inverse=True)
234
238
  _mathscr_v_contribution = _mathscr_v(
235
- scaled_tau, l_inv, Nscoeffs,
239
+ scaled_tau, l_inv,
240
+ Nscoeffs,
236
241
  scaled_s_poly_coeffs[l_uniq],
237
242
  G_collect_0[l_uniq],
238
243
  K_collect_0[l_uniq],
239
- G_inv_collect_0[l_uniq],
240
- mu_arr,
244
+ G_inv_mu_inv[l_uniq],
241
245
  is_antiderivative_wrt_tau,
242
246
  autograd_compatible,
243
247
  )
244
248
 
245
249
  if NFourier == 1:
246
- um = um + _mathscr_v_contribution.T
250
+ um = um + _mathscr_v_contribution
247
251
  elif autograd_compatible:
248
- um = um + np.concatenate((_mathscr_v_contribution.T, np.zeros((NFourier - 1, len(tau), NQuad))))
252
+ um = np.concatenate((um[[0], :, :] + _mathscr_v_contribution, um[1:, :, :]), axis=0)
249
253
  else:
250
- um[0, :, :] += _mathscr_v_contribution.T
254
+ um[0, :, :] += _mathscr_v_contribution
251
255
 
252
- u = np.einsum(
253
- "mti, mp -> itp",
254
- um,
255
- np.cos(np.arange(NFourier)[:, None] * (phi0 - phi)[None, :]),
256
- optimize=True,
256
+ u = np.sum(
257
+ um[:, :, :, None]
258
+ * np.cos(np.arange(NFourier)[:, None] * (phi0 - phi)[None, :])[:, None, None, :],
259
+ axis=0,
257
260
  )
261
+
262
+ outputs = (rescale_factor * np.squeeze(u),)
263
+ output_is_tuple = False
258
264
 
259
265
  if return_Fourier_error:
260
266
  exponent = np.concatenate(
@@ -264,7 +270,14 @@ def _assemble_intensity_and_fluxes(
264
270
  ],
265
271
  axis=-1,
266
272
  )
267
- if is_antiderivative_wrt_tau and there_is_beam_source:
273
+
274
+ direct_beam_contribution = 0
275
+ if is_antiderivative_wrt_tau:
276
+ if there_is_beam_source:
277
+ direct_beam_contribution = (
278
+ B_collect.T[-1, :, :] / (-scale_tau / mu0)[None, :]
279
+ )[:, l] * np.exp(-scaled_tau / mu0)[None, :]
280
+
268
281
  ulast = (
269
282
  np.einsum(
270
283
  "tij, tj -> it",
@@ -272,25 +285,21 @@ def _assemble_intensity_and_fluxes(
272
285
  np.exp(exponent) / (scale_tau[:, None] * K_collect[-1, :, :])[l, :],
273
286
  optimize=True,
274
287
  )
275
- + (B_collect.T[-1, :, :] / (-scale_tau / mu0)[None, :])[:, l]
276
- * np.exp(-scaled_tau / mu0)[None, :]
288
+ + direct_beam_contribution
277
289
  )
278
- elif is_antiderivative_wrt_tau and not there_is_beam_source:
279
- ulast = np.einsum(
280
- "tij, tj -> it",
281
- GC_collect[-1, l, :, :],
282
- np.exp(exponent) / (scale_tau[:, None] * K_collect[-1, :, :])[l, :],
283
- optimize=True,
284
- )
285
- elif not is_antiderivative_wrt_tau and not there_is_beam_source:
286
- ulast = np.einsum(
287
- "tij, tj -> it", GC_collect[-1, l, :, :], np.exp(exponent), optimize=True
288
- ) + B_collect.T[-1, :, l] * np.exp(-scaled_tau / mu0)[None, :]
289
290
  else:
290
- ulast = np.einsum(
291
- "tij, tj -> it", GC_collect[-1, l, :, :], np.exp(exponent), optimize=True
291
+ if there_is_beam_source:
292
+ direct_beam_contribution = (
293
+ B_collect.T[-1, :, l] * np.exp(-scaled_tau / mu0)[None, :]
294
+ )
295
+
296
+ ulast = (
297
+ np.einsum(
298
+ "tij, tj -> it", GC_collect[-1, l, :, :], np.exp(exponent), optimize=True
299
+ )
300
+ + direct_beam_contribution
292
301
  )
293
-
302
+
294
303
  u_abs = np.abs(u)
295
304
  Fourier_error = np.max(
296
305
  np.divide(
@@ -299,19 +308,25 @@ def _assemble_intensity_and_fluxes(
299
308
  ),
300
309
  u_abs,
301
310
  out=np.zeros_like(u_abs),
302
- where=u_abs > 1e-15,
311
+ where=u_abs > 1e-8,
303
312
  )
304
313
  )
305
314
 
306
- if return_tau_arr:
307
- return rescale_factor * np.squeeze(u), Fourier_error, tau_arr
308
- else:
309
- return rescale_factor * np.squeeze(u), Fourier_error
310
- else:
311
- if return_tau_arr:
312
- return rescale_factor * np.squeeze(u), tau_arr
313
- else:
314
- return rescale_factor * np.squeeze(u)
315
+ outputs += (Fourier_error,)
316
+ output_is_tuple = True
317
+
318
+ if return_tau_arr:
319
+ outputs += (tau_arr,)
320
+ output_is_tuple = True
321
+
322
+ if _return_l: # This indicates to the function that NT corrections are activated
323
+ output_is_tuple = True
324
+ outputs += (l,)
325
+ if there_is_iso_source:
326
+ # The TMS correction will use these
327
+ outputs += (l_uniq, l_inv)
328
+
329
+ return outputs if output_is_tuple else outputs[0]
315
330
  # --------------------------------------------------------------------------------------------------------------------------
316
331
 
317
332
  # Construct u0
@@ -332,7 +347,7 @@ def _assemble_intensity_and_fluxes(
332
347
 
333
348
  # Atmospheric layer indices
334
349
  l = np.argmax(tau[:, None] <= tau_arr[None, :], axis=1)
335
- scaled_tau_arr_l = scaled_tau_arr_with_0[l + 1]
350
+ scaled_tau_arr_l = scaled_tau_arr_with_0[1:][l]
336
351
  scaled_tau_arr_lm1 = scaled_tau_arr_with_0[l]
337
352
 
338
353
  # Delta-M scaling
@@ -363,7 +378,13 @@ def _assemble_intensity_and_fluxes(
363
378
  axis=-1,
364
379
  )
365
380
 
381
+ direct_beam_contribution = 0
366
382
  if is_antiderivative_wrt_tau and there_is_beam_source:
383
+ if there_is_beam_source:
384
+ direct_beam_contribution = (B_collect_0.T / (-scale_tau / mu0)[None, :])[
385
+ :, l
386
+ ] * np.exp(-scaled_tau / mu0)[None, :]
387
+
367
388
  u0 = (
368
389
  np.einsum(
369
390
  "tij, tj -> it",
@@ -371,48 +392,46 @@ def _assemble_intensity_and_fluxes(
371
392
  np.exp(exponent) / (scale_tau[:, None] * K_collect_0)[l, :],
372
393
  optimize=True,
373
394
  )
374
- + (B_collect_0.T / (-scale_tau / mu0)[None, :])[:, l]
375
- * np.exp(-scaled_tau / mu0)[None, :]
376
- )
377
- elif is_antiderivative_wrt_tau and not there_is_beam_source:
378
- u0 = np.einsum(
379
- "tij, tj -> it",
380
- GC_collect_0[l, :, :],
381
- np.exp(exponent) / (scale_tau[:, None] * K_collect_0)[l, :],
382
- optimize=True,
395
+ + direct_beam_contribution
383
396
  )
384
- elif not is_antiderivative_wrt_tau and there_is_beam_source:
385
- u0 = np.einsum(
386
- "tij, tj -> it", GC_collect_0[l, :, :], np.exp(exponent), optimize=True
387
- ) + B_collect_0.T[:, l] * np.exp(-scaled_tau / mu0)[None, :]
388
397
  else:
389
- u0 = np.einsum(
390
- "tij, tj -> it", GC_collect_0[l, :, :], np.exp(exponent), optimize=True
398
+ if there_is_beam_source:
399
+ direct_beam_contribution = (
400
+ B_collect_0.T[:, l] * np.exp(-scaled_tau / mu0)[None, :]
401
+ )
402
+
403
+ u0 = (
404
+ np.einsum(
405
+ "tij, tj -> it", GC_collect_0[l, :, :], np.exp(exponent), optimize=True
406
+ )
407
+ + direct_beam_contribution
391
408
  )
392
-
409
+
393
410
  # Contribution from particular solution for isotropic internal sources
394
411
  if there_is_iso_source:
395
412
  l_uniq, l_inv = np.unique(l, return_inverse=True)
396
413
  _mathscr_v_contribution = _mathscr_v(
397
- scaled_tau, l_inv, Nscoeffs,
414
+ scaled_tau, l_inv,
415
+ Nscoeffs,
398
416
  scaled_s_poly_coeffs[l_uniq],
399
417
  G_collect_0[l_uniq],
400
418
  K_collect_0[l_uniq],
401
- G_inv_collect_0[l_uniq],
402
- mu_arr,
419
+ G_inv_mu_inv[l_uniq],
403
420
  is_antiderivative_wrt_tau,
404
421
  autograd_compatible,
405
422
  )
406
423
  u0 = u0 + _mathscr_v_contribution
407
424
 
408
- if return_tau_arr and _return_act_dscale_for_reclass:
409
- return rescale_factor * np.squeeze(u0), tau_arr, act_dscale_reclassification
410
- elif return_tau_arr and not _return_act_dscale_for_reclass:
411
- return rescale_factor * np.squeeze(u0), tau_arr
412
- elif not return_tau_arr and _return_act_dscale_for_reclass:
413
- return rescale_factor * np.squeeze(u0), act_dscale_reclassification
414
- else:
415
- return rescale_factor * np.squeeze(u0)
425
+ outputs = (rescale_factor * np.squeeze(u0),)
426
+ outputs_len = 1
427
+ if return_tau_arr:
428
+ outputs += (tau_arr,)
429
+ outputs_len += 1
430
+ if _return_act_dscale_for_reclass:
431
+ outputs += (act_dscale_reclassification,)
432
+ outputs_len += 1
433
+ return outputs[0] if outputs_len == 1 else outputs
434
+
416
435
  # --------------------------------------------------------------------------------------------------------------------------
417
436
 
418
437
  # Construct the flux functions (refer to section 3.8 of the Comprehensive Documentation)
@@ -437,7 +456,7 @@ def _assemble_intensity_and_fluxes(
437
456
 
438
457
  # Atmospheric layer indices
439
458
  l = np.argmax(tau[:, None] <= tau_arr[None, :], axis=1)
440
- scaled_tau_arr_l = scaled_tau_arr_with_0[l + 1]
459
+ scaled_tau_arr_l = scaled_tau_arr_with_0[1:][l]
441
460
  scaled_tau_arr_lm1 = scaled_tau_arr_with_0[l]
442
461
 
443
462
  # Delta-M scaling
@@ -451,27 +470,18 @@ def _assemble_intensity_and_fluxes(
451
470
  if there_is_iso_source:
452
471
  l_uniq, l_inv = np.unique(l, return_inverse=True)
453
472
  _mathscr_v_contribution = _mathscr_v(
454
- scaled_tau, l_inv, Nscoeffs,
473
+ scaled_tau, l_inv,
474
+ Nscoeffs,
455
475
  scaled_s_poly_coeffs[l_uniq],
456
476
  G_collect_0[l_uniq, :N, :],
457
477
  K_collect_0[l_uniq],
458
- G_inv_collect_0[l_uniq],
459
- mu_arr,
478
+ G_inv_mu_inv[l_uniq],
460
479
  is_antiderivative_wrt_tau,
461
480
  autograd_compatible,
462
481
  )
463
482
  else:
464
483
  _mathscr_v_contribution = 0
465
484
 
466
- if there_is_beam_source and is_antiderivative_wrt_tau:
467
- direct_beam_contribution = (
468
- (B_pos / (-scale_tau / mu0)[None, :])[:, l] * np.exp(-scaled_tau / mu0)[None, :]
469
- )
470
- elif there_is_beam_source and not is_antiderivative_wrt_tau:
471
- direct_beam_contribution = B_pos[:, l] * np.exp(-scaled_tau / mu0)[None, :]
472
- else:
473
- direct_beam_contribution = 0
474
-
475
485
  exponent = np.concatenate(
476
486
  [
477
487
  K_collect_0[l, :N] * (scaled_tau - scaled_tau_arr_lm1)[:, None],
@@ -480,7 +490,13 @@ def _assemble_intensity_and_fluxes(
480
490
  axis=-1,
481
491
  )
482
492
 
493
+ direct_beam_contribution = 0
483
494
  if is_antiderivative_wrt_tau:
495
+ if there_is_beam_source:
496
+ direct_beam_contribution = (
497
+ (B_pos / (-scale_tau / mu0)[None, :])[:, l] * np.exp(-scaled_tau / mu0)[None, :]
498
+ )
499
+
484
500
  u0_pos = (
485
501
  np.einsum(
486
502
  "tij, tj -> it",
@@ -491,11 +507,15 @@ def _assemble_intensity_and_fluxes(
491
507
  + _mathscr_v_contribution
492
508
  )
493
509
  else:
510
+ if there_is_beam_source:
511
+ direct_beam_contribution = B_pos[:, l] * np.exp(-scaled_tau / mu0)[None, :]
512
+
494
513
  u0_pos = (
495
514
  np.einsum("tij, tj -> it", GC_pos[l, :, :], np.exp(exponent))
496
515
  + direct_beam_contribution
497
516
  + _mathscr_v_contribution
498
517
  )
518
+
499
519
  flux = 2 * pi * (mu_arr_pos * W) @ u0_pos
500
520
 
501
521
  if return_tau_arr:
@@ -518,7 +538,7 @@ def _assemble_intensity_and_fluxes(
518
538
 
519
539
  # Atmospheric layer indices
520
540
  l = np.argmax(tau[:, None] <= tau_arr[None, :], axis=1)
521
- scaled_tau_arr_l = scaled_tau_arr_with_0[l + 1]
541
+ scaled_tau_arr_l = scaled_tau_arr_with_0[1:][l]
522
542
  scaled_tau_arr_lm1 = scaled_tau_arr_with_0[l]
523
543
 
524
544
  # Delta-M scaling
@@ -532,31 +552,18 @@ def _assemble_intensity_and_fluxes(
532
552
  if there_is_iso_source:
533
553
  l_uniq, l_inv = np.unique(l, return_inverse=True)
534
554
  _mathscr_v_contribution = _mathscr_v(
535
- scaled_tau, l_inv, Nscoeffs,
555
+ scaled_tau, l_inv,
556
+ Nscoeffs,
536
557
  scaled_s_poly_coeffs[l_uniq],
537
558
  G_collect_0[l_uniq, N:, :],
538
559
  K_collect_0[l_uniq],
539
- G_inv_collect_0[l_uniq],
540
- mu_arr,
560
+ G_inv_mu_inv[l_uniq],
541
561
  is_antiderivative_wrt_tau,
542
562
  autograd_compatible,
543
563
  )
544
564
  else:
545
565
  _mathscr_v_contribution = 0
546
-
547
- if there_is_beam_source and is_antiderivative_wrt_tau:
548
- direct_beam_contribution = (B_neg / (-scale_tau / mu0)[None, l])[:, l] * np.exp(-scaled_tau / mu0)[None, :]
549
- direct_beam = I0 * mu0 * np.exp(-tau / mu0) * -mu0
550
- direct_beam_scaled = I0 * mu0 * np.exp(-scaled_tau / mu0) / (-scale_tau / mu0)[None, l]
551
- elif there_is_beam_source and not is_antiderivative_wrt_tau:
552
- direct_beam_contribution = B_neg[:, l] * np.exp(-scaled_tau / mu0)[None, :]
553
- direct_beam = I0 * mu0 * np.exp(-tau / mu0)
554
- direct_beam_scaled = I0 * mu0 * np.exp(-scaled_tau / mu0)
555
- else:
556
- direct_beam_contribution = 0
557
- direct_beam = 0
558
- direct_beam_scaled = 0
559
-
566
+
560
567
  exponent = np.concatenate(
561
568
  [
562
569
  K_collect_0[l, :N] * (scaled_tau - scaled_tau_arr_lm1)[:, None],
@@ -565,18 +572,32 @@ def _assemble_intensity_and_fluxes(
565
572
  axis=-1,
566
573
  )
567
574
 
575
+ direct_beam_contribution = 0
576
+ direct_beam = 0
577
+ direct_beam_scaled = 0
568
578
  if is_antiderivative_wrt_tau:
579
+ if there_is_beam_source:
580
+ direct_beam_contribution = (B_neg / (-scale_tau / mu0)[None, l])[:, l] * np.exp(-scaled_tau / mu0)[None, :]
581
+ direct_beam = I0 * mu0 * np.exp(-tau / mu0) * -mu0
582
+ direct_beam_scaled = I0 * mu0 * np.exp(-scaled_tau / mu0) / (-scale_tau / mu0)[None, l]
583
+
569
584
  u0_neg = (
570
585
  np.einsum("tij, tj -> it", GC_neg[l, :, :], np.exp(exponent) / (scale_tau[:, None] * K_collect_0)[l, :])
571
586
  + direct_beam_contribution
572
587
  + _mathscr_v_contribution
573
588
  )
574
589
  else:
590
+ if there_is_beam_source:
591
+ direct_beam_contribution = B_neg[:, l] * np.exp(-scaled_tau / mu0)[None, :]
592
+ direct_beam = I0 * mu0 * np.exp(-tau / mu0)
593
+ direct_beam_scaled = I0 * mu0 * np.exp(-scaled_tau / mu0)
594
+
575
595
  u0_neg = (
576
596
  np.einsum("tij, tj -> it", GC_neg[l, :, :], np.exp(exponent))
577
597
  + direct_beam_contribution
578
598
  + _mathscr_v_contribution
579
599
  )
600
+
580
601
  diffuse_flux = 2 * pi * (mu_arr_pos * W) @ u0_neg + direct_beam_scaled - direct_beam
581
602
 
582
603
  if return_tau_arr: