PythonicDISORT 1.6__py3-none-any.whl → 1.7__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.
- PythonicDISORT/_assemble_intensity_and_fluxes.py +18 -14
- PythonicDISORT/_solve_for_coeffs.py +62 -69
- PythonicDISORT/_solve_for_gen_and_part_sols.py +41 -61
- PythonicDISORT/pydisort.py +85 -66
- PythonicDISORT/subroutines.py +35 -20
- {pythonicdisort-1.6.dist-info → pythonicdisort-1.7.dist-info}/METADATA +1 -1
- pythonicdisort-1.7.dist-info/RECORD +11 -0
- {pythonicdisort-1.6.dist-info → pythonicdisort-1.7.dist-info}/WHEEL +1 -1
- pythonicdisort-1.6.dist-info/RECORD +0 -11
- {pythonicdisort-1.6.dist-info → pythonicdisort-1.7.dist-info}/LICENSE.md +0 -0
- {pythonicdisort-1.6.dist-info → pythonicdisort-1.7.dist-info}/top_level.txt +0 -0
|
@@ -182,12 +182,12 @@ def _assemble_intensity_and_fluxes(
|
|
|
182
182
|
raise ValueError("tau input outside the tau range specified for the atmosphere (check `tau_arr`).")
|
|
183
183
|
|
|
184
184
|
# Atmospheric layer indices
|
|
185
|
-
l = np.
|
|
185
|
+
l = np.searchsorted(tau_arr, tau, side="left")
|
|
186
186
|
scaled_tau_arr_l = scaled_tau_arr_with_0[1:][l]
|
|
187
187
|
scaled_tau_arr_lm1 = scaled_tau_arr_with_0[l]
|
|
188
188
|
|
|
189
189
|
# Delta-M scaling
|
|
190
|
-
if np.any(scale_tau !=
|
|
190
|
+
if np.any(scale_tau != 1):
|
|
191
191
|
tau_dist_from_top = tau_arr[l] - tau
|
|
192
192
|
scaled_tau_dist_from_top = tau_dist_from_top * scale_tau[l]
|
|
193
193
|
scaled_tau = scaled_tau_arr_l - scaled_tau_dist_from_top
|
|
@@ -236,7 +236,8 @@ def _assemble_intensity_and_fluxes(
|
|
|
236
236
|
if there_is_iso_source:
|
|
237
237
|
l_uniq, l_inv = np.unique(l, return_inverse=True)
|
|
238
238
|
_mathscr_v_contribution = _mathscr_v(
|
|
239
|
-
scaled_tau, scale_tau,
|
|
239
|
+
scaled_tau, scale_tau[l_uniq], # `scale_tau` is only required to determine the antiderivative
|
|
240
|
+
l_inv,
|
|
240
241
|
Nscoeffs,
|
|
241
242
|
scaled_s_poly_coeffs[l_uniq],
|
|
242
243
|
G_collect_0[l_uniq],
|
|
@@ -339,12 +340,12 @@ def _assemble_intensity_and_fluxes(
|
|
|
339
340
|
raise ValueError("tau input outside the tau range given for the atmosphere (check `tau_arr`).")
|
|
340
341
|
|
|
341
342
|
# Atmospheric layer indices
|
|
342
|
-
l = np.
|
|
343
|
+
l = np.searchsorted(tau_arr, tau, side="left")
|
|
343
344
|
scaled_tau_arr_l = scaled_tau_arr_with_0[1:][l]
|
|
344
345
|
scaled_tau_arr_lm1 = scaled_tau_arr_with_0[l]
|
|
345
346
|
|
|
346
347
|
# Delta-M scaling
|
|
347
|
-
if np.any(scale_tau !=
|
|
348
|
+
if np.any(scale_tau != 1):
|
|
348
349
|
tau_dist_from_top = tau_arr[l] - tau
|
|
349
350
|
scaled_tau_dist_from_top = tau_dist_from_top * scale_tau[l]
|
|
350
351
|
scaled_tau = scaled_tau_arr_l - scaled_tau_dist_from_top
|
|
@@ -404,7 +405,8 @@ def _assemble_intensity_and_fluxes(
|
|
|
404
405
|
if there_is_iso_source:
|
|
405
406
|
l_uniq, l_inv = np.unique(l, return_inverse=True)
|
|
406
407
|
_mathscr_v_contribution = _mathscr_v(
|
|
407
|
-
scaled_tau, scale_tau,
|
|
408
|
+
scaled_tau, scale_tau[l_uniq], # `scale_tau` is only required to determine the antiderivative
|
|
409
|
+
l_inv,
|
|
408
410
|
Nscoeffs,
|
|
409
411
|
scaled_s_poly_coeffs[l_uniq],
|
|
410
412
|
G_collect_0[l_uniq],
|
|
@@ -448,12 +450,12 @@ def _assemble_intensity_and_fluxes(
|
|
|
448
450
|
raise ValueError("tau input outside the tau range given for the atmosphere (check `tau_arr`).")
|
|
449
451
|
|
|
450
452
|
# Atmospheric layer indices
|
|
451
|
-
l = np.
|
|
453
|
+
l = np.searchsorted(tau_arr, tau, side="left")
|
|
452
454
|
scaled_tau_arr_l = scaled_tau_arr_with_0[1:][l]
|
|
453
455
|
scaled_tau_arr_lm1 = scaled_tau_arr_with_0[l]
|
|
454
456
|
|
|
455
457
|
# Delta-M scaling
|
|
456
|
-
if np.any(scale_tau !=
|
|
458
|
+
if np.any(scale_tau != 1):
|
|
457
459
|
tau_dist_from_top = tau_arr[l] - tau
|
|
458
460
|
scaled_tau_dist_from_top = tau_dist_from_top * scale_tau[l]
|
|
459
461
|
scaled_tau = scaled_tau_arr_l - scaled_tau_dist_from_top
|
|
@@ -463,7 +465,8 @@ def _assemble_intensity_and_fluxes(
|
|
|
463
465
|
if there_is_iso_source:
|
|
464
466
|
l_uniq, l_inv = np.unique(l, return_inverse=True)
|
|
465
467
|
_mathscr_v_contribution = _mathscr_v(
|
|
466
|
-
scaled_tau, scale_tau,
|
|
468
|
+
scaled_tau, scale_tau[l_uniq], # `scale_tau` is only required to determine the antiderivative
|
|
469
|
+
l_inv,
|
|
467
470
|
Nscoeffs,
|
|
468
471
|
scaled_s_poly_coeffs[l_uniq],
|
|
469
472
|
G_collect_0[l_uniq, :N, :],
|
|
@@ -530,12 +533,12 @@ def _assemble_intensity_and_fluxes(
|
|
|
530
533
|
raise ValueError("tau input outside the tau range given for the atmosphere (check `tau_arr`).")
|
|
531
534
|
|
|
532
535
|
# Atmospheric layer indices
|
|
533
|
-
l = np.
|
|
536
|
+
l = np.searchsorted(tau_arr, tau, side="left")
|
|
534
537
|
scaled_tau_arr_l = scaled_tau_arr_with_0[1:][l]
|
|
535
538
|
scaled_tau_arr_lm1 = scaled_tau_arr_with_0[l]
|
|
536
539
|
|
|
537
540
|
# Delta-M scaling
|
|
538
|
-
if np.any(scale_tau !=
|
|
541
|
+
if np.any(scale_tau != 1):
|
|
539
542
|
tau_dist_from_top = tau_arr[l] - tau
|
|
540
543
|
scaled_tau_dist_from_top = tau_dist_from_top * scale_tau[l]
|
|
541
544
|
scaled_tau = scaled_tau_arr_l - scaled_tau_dist_from_top
|
|
@@ -545,7 +548,8 @@ def _assemble_intensity_and_fluxes(
|
|
|
545
548
|
if there_is_iso_source:
|
|
546
549
|
l_uniq, l_inv = np.unique(l, return_inverse=True)
|
|
547
550
|
_mathscr_v_contribution = _mathscr_v(
|
|
548
|
-
scaled_tau, scale_tau,
|
|
551
|
+
scaled_tau, scale_tau[l_uniq], # `scale_tau` is only required to determine the antiderivative
|
|
552
|
+
l_inv,
|
|
549
553
|
Nscoeffs,
|
|
550
554
|
scaled_s_poly_coeffs[l_uniq],
|
|
551
555
|
G_collect_0[l_uniq, N:, :],
|
|
@@ -570,9 +574,9 @@ def _assemble_intensity_and_fluxes(
|
|
|
570
574
|
direct_beam_scaled = 0
|
|
571
575
|
if is_antiderivative_wrt_tau:
|
|
572
576
|
if there_is_beam_source:
|
|
573
|
-
direct_beam_contribution = (B_neg / (-scale_tau / mu0)[None,
|
|
577
|
+
direct_beam_contribution = (B_neg / (-scale_tau / mu0)[None, :])[:, l] * np.exp(-scaled_tau / mu0)[None, :]
|
|
574
578
|
direct_beam = I0 * mu0 * np.exp(-tau / mu0) * -mu0
|
|
575
|
-
direct_beam_scaled = I0 * mu0 * np.exp(-scaled_tau / mu0) / (-scale_tau / mu0)
|
|
579
|
+
direct_beam_scaled = I0 * mu0 * np.exp(-scaled_tau / mu0) / (-scale_tau[l] / mu0)
|
|
576
580
|
|
|
577
581
|
u0_neg = (
|
|
578
582
|
np.einsum("tij, tj -> it", GC_neg[l, :, :], np.exp(exponent) / (scale_tau[:, None] * K_collect_0)[l, :])
|
|
@@ -79,7 +79,8 @@ def _solve_for_coeffs(
|
|
|
79
79
|
use_banded_solver = (NLayers >= use_banded_solver_NLayers)
|
|
80
80
|
dim = NLayers * NQuad
|
|
81
81
|
C_collect = np.empty((NFourier, dim))
|
|
82
|
-
|
|
82
|
+
thickness = np.diff(scaled_tau_arr_with_0) # (NLayers,)
|
|
83
|
+
|
|
83
84
|
RHS = np.empty(dim)
|
|
84
85
|
i_bot = dim - N # Equals `j_bot_right`
|
|
85
86
|
j_bot_left = dim - NQuad
|
|
@@ -89,23 +90,47 @@ def _solve_for_coeffs(
|
|
|
89
90
|
LHS_dof = np.zeros((6 * N - 1, dim), order="F") # Diagonal ordered form
|
|
90
91
|
s0, s1 = LHS_dof.strides
|
|
91
92
|
col_stride = s1 - s0 # Stride when moving +1 col and -1 row
|
|
92
|
-
|
|
93
|
+
|
|
93
94
|
def _view_slanted(i0, j0, nrows):
|
|
94
|
-
"""
|
|
95
|
-
|
|
96
|
-
|
|
97
|
-
|
|
95
|
+
"""Single slanted view into `LHS_dof` anchored at conceptual LHS coords (i0, j0).
|
|
96
|
+
Returns a 2D view of shape (nrows, N). Used for top/bottom boundary-condition rows.
|
|
97
|
+
|
|
98
|
+
Uses `numpy.lib.stride_tricks.as_strided` (zero-copy but unsafe), which is why
|
|
99
|
+
this helper is not exposed in `PythonicDISORT.subroutines`.
|
|
98
100
|
See https://numpy.org/devdocs/reference/generated/numpy.lib.stride_tricks.as_strided.html.
|
|
99
|
-
|
|
101
|
+
|
|
100
102
|
Assume the following exist in the outer scope:
|
|
101
|
-
`LHS_dof`, `Nsupsubdiags`, `s0`, `col_stride
|
|
103
|
+
`LHS_dof`, `Nsupsubdiags`, `N`, `s0`, `col_stride`.
|
|
102
104
|
"""
|
|
103
105
|
r0 = Nsupsubdiags + i0 - j0
|
|
104
106
|
base = LHS_dof[r0 : r0 + nrows, j0 : j0 + N]
|
|
105
107
|
return as_strided(base, shape=(nrows, N), strides=(s0, col_stride))
|
|
106
|
-
|
|
108
|
+
|
|
109
|
+
def _view_slanted_batched(i0, j0, nrows):
|
|
110
|
+
"""Batched (NLayers-1, nrows, N) slanted view across interlayer continuity rows.
|
|
111
|
+
Anchored at conceptual LHS coords (i0, j0) for layer 0; layer l advances both
|
|
112
|
+
i and j by `NQuad` along the diagonal.
|
|
113
|
+
|
|
114
|
+
Same `(i0, j0)` convention as `_view_slanted` (slant arithmetic done internally).
|
|
115
|
+
Assume the following exist in the outer scope:
|
|
116
|
+
`LHS_dof`, `Nsupsubdiags`, `N`, `NQuad`, `NLayers`, `s0`, `s1`, `col_stride`.
|
|
117
|
+
"""
|
|
118
|
+
r0 = Nsupsubdiags + i0 - j0 # matches `_view_slanted`'s internal slant
|
|
119
|
+
base = LHS_dof[r0 : r0 + nrows, j0 : j0 + N + (NLayers - 2) * NQuad]
|
|
120
|
+
return as_strided(base, shape=(NLayers - 1, nrows, N),
|
|
121
|
+
strides=(NQuad * s1, s0, col_stride))
|
|
122
|
+
|
|
107
123
|
else:
|
|
108
124
|
LHS = np.zeros((dim, dim), order="F")
|
|
125
|
+
sR, sC = LHS.strides
|
|
126
|
+
def _dense_slab(r0, c0, nrows, ncols):
|
|
127
|
+
"""Batched (NLayers-1, nrows, ncols) view across interlayer blocks of `LHS`.
|
|
128
|
+
Anchored at LHS coords (r0, c0) for layer 0; layer l advances both row and column
|
|
129
|
+
anchors by `NQuad` along the diagonal. Each slab has shape (nrows, ncols)."""
|
|
130
|
+
base = LHS[r0 : r0 + nrows + (NLayers - 2) * NQuad,
|
|
131
|
+
c0 : c0 + ncols + (NLayers - 2) * NQuad]
|
|
132
|
+
return as_strided(base, shape=(NLayers - 1, nrows, ncols),
|
|
133
|
+
strides=(NQuad * (sR + sC), sR, sC))
|
|
109
134
|
|
|
110
135
|
# The following loops can easily be parallelized, but the speed-up is unlikely to be worth the overhead
|
|
111
136
|
for m in range(NFourier):
|
|
@@ -262,7 +287,7 @@ def _solve_for_coeffs(
|
|
|
262
287
|
G_0_nn = G_0[N:, :N]
|
|
263
288
|
G_0_np = G_0[N:, N:]
|
|
264
289
|
E_01 = np.exp(K_collect_m[0, :N] * scaled_tau_arr_with_0[1])[None, :]
|
|
265
|
-
|
|
290
|
+
|
|
266
291
|
G_L = G_collect_m[-1]
|
|
267
292
|
G_L_pn = G_L[:N, :N]
|
|
268
293
|
G_L_nn = G_L[N:, :N]
|
|
@@ -271,6 +296,9 @@ def _solve_for_coeffs(
|
|
|
271
296
|
E_Lm1L = np.exp(
|
|
272
297
|
K_collect_m[-1, :N] * (scaled_tau_arr_with_0[-1] - scaled_tau_arr_with_0[-2])
|
|
273
298
|
)[None, :]
|
|
299
|
+
# Per-layer thickness exponentials, computed once per Fourier mode and reused below.
|
|
300
|
+
# If looping over `l in range(NLayers)`, the line below would be `E_pos_full[l] = exp(-thickness[l] * K_pos[l])`
|
|
301
|
+
E_pos_full = np.exp(-thickness[:, None] * K_collect_m[:, N:]) # (NLayers, N)
|
|
274
302
|
|
|
275
303
|
################################## Use `sc.linalg.solve_banded` ##################################
|
|
276
304
|
if use_banded_solver:
|
|
@@ -292,35 +320,18 @@ def _solve_for_coeffs(
|
|
|
292
320
|
_view_slanted(i_bot, j_bot_left, N)[:] = G_L_pn * E_Lm1L
|
|
293
321
|
_view_slanted(i_bot, i_bot, N)[:] = G_L_pp
|
|
294
322
|
|
|
295
|
-
# ---------------- Interlayer / continuity BCs ----------------
|
|
296
|
-
|
|
297
|
-
|
|
298
|
-
|
|
299
|
-
|
|
300
|
-
G_l_ap = G_l[:, N:] # (2N,N)
|
|
301
|
-
|
|
302
|
-
G_lp1 = G_collect_m[l + 1]
|
|
303
|
-
G_lp1_an = G_lp1[:, :N] # (2N,N)
|
|
304
|
-
G_lp1_pp = G_lp1[:N, N:] # (N,N)
|
|
305
|
-
G_lp1_np = G_lp1[N:, N:] # (N,N)
|
|
306
|
-
|
|
307
|
-
tau_lm1 = scaled_tau_arr_with_0[l]
|
|
308
|
-
tau_l = scaled_tau_arr_with_0[l + 1]
|
|
309
|
-
tau_lp1 = scaled_tau_arr_with_0[l + 2]
|
|
310
|
-
|
|
311
|
-
E_lm1l = np.exp(K_collect_m[l, N:] * (tau_lm1 - tau_l))[None, :] # (N,)
|
|
312
|
-
E_llp1 = np.exp(K_collect_m[l + 1, N:] * (tau_l - tau_lp1))[None, :] # (N,)
|
|
323
|
+
# ---------------- Interlayer / continuity BCs (batched across layers) ----------------
|
|
324
|
+
E_lm1l = E_pos_full[:-1, None, :] # (NLayers-1, 1, N)
|
|
325
|
+
E_llp1 = E_pos_full[1:, None, :] # (NLayers-1, 1, N)
|
|
326
|
+
G_lo = G_collect_m[:-1] # (NLayers-1, NQuad, NQuad)
|
|
327
|
+
G_hi = G_collect_m[1:] # (NLayers-1, NQuad, NQuad)
|
|
313
328
|
|
|
314
|
-
|
|
315
|
-
|
|
316
|
-
|
|
317
|
-
|
|
318
|
-
|
|
319
|
-
|
|
320
|
-
|
|
321
|
-
_view_slanted(start_row, start_col + NQuad, NQuad)[:] = -G_lp1_an
|
|
322
|
-
_view_slanted(start_row, start_col + 3 * N, N)[:] = -G_lp1_pp * E_llp1
|
|
323
|
-
_view_slanted(start_row + N, start_col + 3 * N, N)[:] = -G_lp1_np * E_llp1
|
|
329
|
+
_view_slanted_batched(N, 0, N)[:] = G_lo[:, :N, :N] * E_lm1l
|
|
330
|
+
_view_slanted_batched(2 * N, 0, N)[:] = G_lo[:, N:, :N] * E_lm1l
|
|
331
|
+
_view_slanted_batched(N, N, NQuad)[:] = G_lo[:, :, N:]
|
|
332
|
+
_view_slanted_batched(N, NQuad, NQuad)[:] = -G_hi[:, :, :N]
|
|
333
|
+
_view_slanted_batched(N, 3 * N, N)[:] = -G_hi[:, :N, N:] * E_llp1
|
|
334
|
+
_view_slanted_batched(2 * N, 3 * N, N)[:] = -G_hi[:, N:, N:] * E_llp1
|
|
324
335
|
|
|
325
336
|
# Solve the system
|
|
326
337
|
C_m = sc.linalg.solve_banded(
|
|
@@ -329,7 +340,7 @@ def _solve_for_coeffs(
|
|
|
329
340
|
RHS,
|
|
330
341
|
False,
|
|
331
342
|
True,
|
|
332
|
-
|
|
343
|
+
True,
|
|
333
344
|
)
|
|
334
345
|
|
|
335
346
|
##################################################################################################
|
|
@@ -348,36 +359,18 @@ def _solve_for_coeffs(
|
|
|
348
359
|
LHS[-N:, -NQuad : -N] = G_L_pn * E_Lm1L
|
|
349
360
|
LHS[-N:, -N:] = G_L_pp
|
|
350
361
|
|
|
351
|
-
# ---------------- Interlayer / continuity BCs ----------------
|
|
352
|
-
|
|
353
|
-
|
|
354
|
-
|
|
355
|
-
|
|
356
|
-
|
|
357
|
-
|
|
358
|
-
|
|
359
|
-
|
|
360
|
-
|
|
361
|
-
|
|
362
|
-
|
|
363
|
-
scaled_tau_arr_lm1 = scaled_tau_arr_with_0[l]
|
|
364
|
-
scaled_tau_arr_l = scaled_tau_arr_with_0[l + 1]
|
|
365
|
-
scaled_tau_arr_lp1 = scaled_tau_arr_with_0[l + 2]
|
|
366
|
-
|
|
367
|
-
# Postive eigenvalues
|
|
368
|
-
E_lm1l = np.exp(K_collect_m[l, N:] * (scaled_tau_arr_lm1 - scaled_tau_arr_l))
|
|
369
|
-
E_llp1 = np.exp(K_collect_m[l + 1, N:] * (scaled_tau_arr_l - scaled_tau_arr_lp1))
|
|
370
|
-
|
|
371
|
-
start_row = N + l * NQuad
|
|
372
|
-
start_col = l * NQuad
|
|
373
|
-
|
|
374
|
-
LHS[start_row : N + start_row, start_col : N + start_col] = G_l_pn * E_lm1l
|
|
375
|
-
LHS[N + start_row : NQuad + start_row, start_col : N + start_col] = G_l_nn * E_lm1l
|
|
376
|
-
LHS[start_row : NQuad + start_row, N + start_col : NQuad + start_col] = G_l_ap
|
|
377
|
-
|
|
378
|
-
LHS[start_row : NQuad + start_row, NQuad + start_col : 3 * N + start_col] = -G_lp1_an
|
|
379
|
-
LHS[start_row : N + start_row, 3 * N + start_col : 4 * N + start_col] = -G_lp1_pp * E_llp1
|
|
380
|
-
LHS[N + start_row : NQuad + start_row, 3 * N + start_col : 4 * N + start_col] = -G_lp1_np * E_llp1
|
|
362
|
+
# ---------------- Interlayer / continuity BCs (batched across layers) ----------------
|
|
363
|
+
E_lm1l = E_pos_full[:-1, None, :] # (NLayers-1, 1, N)
|
|
364
|
+
E_llp1 = E_pos_full[1:, None, :] # (NLayers-1, 1, N)
|
|
365
|
+
G_lo = G_collect_m[:-1]
|
|
366
|
+
G_hi = G_collect_m[1:]
|
|
367
|
+
|
|
368
|
+
_dense_slab(N, 0, N, N)[:] = G_lo[:, :N, :N] * E_lm1l
|
|
369
|
+
_dense_slab(2 * N, 0, N, N)[:] = G_lo[:, N:, :N] * E_lm1l
|
|
370
|
+
_dense_slab(N, N, NQuad, N)[:] = G_lo[:, :, N:]
|
|
371
|
+
_dense_slab(N, NQuad, NQuad, N)[:] = -G_hi[:, :, :N]
|
|
372
|
+
_dense_slab(N, 3 * N, N, N)[:] = -G_hi[:, :N, N:] * E_llp1
|
|
373
|
+
_dense_slab(2 * N, 3 * N, N, N)[:] = -G_hi[:, N:, N:] * E_llp1
|
|
381
374
|
|
|
382
375
|
# Solve the system
|
|
383
376
|
C_m = np.linalg.solve(LHS, RHS)
|
|
@@ -69,13 +69,14 @@ def _solve_for_gen_and_part_sols(
|
|
|
69
69
|
G_trivial.ravel()[N : N + N * stride : stride] = 1
|
|
70
70
|
G_trivial.ravel()[NQuad * N : NQuad * N + N * stride : stride] = 1
|
|
71
71
|
|
|
72
|
+
# Flat (NFourier * NLayers, ...) storage so a single integer index list can address
|
|
73
|
+
# every (m, l) pair across Fourier modes for the cross-mode batched `eig`.
|
|
72
74
|
G_collect = np.empty((NFourier * NLayers, NQuad, NQuad))
|
|
73
75
|
K_collect = np.empty((NFourier * NLayers, NQuad))
|
|
74
76
|
A_arr = np.empty((NFourier * NLayers, NQuad, NQuad))
|
|
75
77
|
no_shortcut_indices = []
|
|
76
78
|
no_shortcut_indices_0 = []
|
|
77
|
-
|
|
78
|
-
|
|
79
|
+
|
|
79
80
|
if there_is_beam_source:
|
|
80
81
|
mu_arr_pos_mu0 = np.append(mu_arr_pos, -mu0)
|
|
81
82
|
|
|
@@ -109,66 +110,45 @@ def _solve_for_gen_and_part_sols(
|
|
|
109
110
|
fac_asso_leg_term_posT = poch[None, :] * asso_leg_term_pos.T
|
|
110
111
|
# --------------------------------------------------------------------------------------------------------------------------
|
|
111
112
|
|
|
112
|
-
#
|
|
113
|
-
# --------------------------------------------------------------------------------------------------------------------------
|
|
114
|
-
|
|
115
|
-
|
|
116
|
-
|
|
117
|
-
|
|
118
|
-
|
|
119
|
-
|
|
120
|
-
|
|
121
|
-
|
|
122
|
-
|
|
123
|
-
|
|
124
|
-
|
|
125
|
-
|
|
126
|
-
|
|
127
|
-
|
|
128
|
-
|
|
113
|
+
# Vectorized over NLayers atmospheric layers
|
|
114
|
+
# --------------------------------------------------------------------------------------------------------------------------
|
|
115
|
+
omega_times_Leg = (scaled_omega_arr[:, None] / 2) * weighted_scaled_Leg_coeffs[:, m:] # (NLayers, NLeg-m)
|
|
116
|
+
no_shortcut = np.any(np.abs(omega_times_Leg) > 1e-8, axis=1) # (NLayers,)
|
|
117
|
+
ns_ind = np.flatnonzero(no_shortcut)
|
|
118
|
+
sh_ind = np.flatnonzero(~no_shortcut)
|
|
119
|
+
ind_base = m * NLayers
|
|
120
|
+
|
|
121
|
+
# Shortcut layers: trivial diagonalization (no multiple scattering)
|
|
122
|
+
if sh_ind.size:
|
|
123
|
+
G_collect[ind_base + sh_ind] = G_trivial
|
|
124
|
+
K_collect[ind_base + sh_ind, :N] = -M_inv
|
|
125
|
+
K_collect[ind_base + sh_ind, N:] = M_inv
|
|
126
|
+
if m_equals_0 and there_is_iso_source:
|
|
127
|
+
G_inv_collect_0[sh_ind] = G_trivial
|
|
128
|
+
|
|
129
|
+
# Non-shortcut layers: full computation, batched over NLayers
|
|
130
|
+
if ns_ind.size:
|
|
131
|
+
D_temp = omega_times_Leg[ns_ind][:, None, :] * fac_asso_leg_term_posT[None, :, :] # (Nns, N, NLeg-m)
|
|
132
|
+
D_pos = D_temp @ asso_leg_term_pos # (Nns, N, N)
|
|
133
|
+
D_neg = D_temp @ asso_leg_term_neg # (Nns, N, N)
|
|
134
|
+
DW = D_pos * W[None, None, :]
|
|
135
|
+
DW.reshape(-1, N * N)[:, ::N + 1] -= 1 # batched diag-sub
|
|
136
|
+
inds = ind_base + ns_ind
|
|
137
|
+
A_arr[inds, N:, N:] = M_inv[None, :, None] * DW # alpha
|
|
138
|
+
A_arr[inds, N:, :N] = M_inv[None, :, None] * D_neg * W[None, None, :] # beta
|
|
139
|
+
no_shortcut_indices.extend(inds.tolist())
|
|
140
|
+
if m_equals_0 and there_is_iso_source:
|
|
141
|
+
no_shortcut_indices_0.extend(ns_ind.tolist())
|
|
142
|
+
if there_is_beam_source:
|
|
143
|
+
X_temp = (
|
|
144
|
+
scalar_fac_asso_leg_term_mu0
|
|
145
|
+
* scaled_omega_arr[ns_ind, None]
|
|
146
|
+
* weighted_scaled_Leg_coeffs[ns_ind, m:]
|
|
147
|
+
)
|
|
148
|
+
# Signs flipped due to the extra minus sign
|
|
149
|
+
X_arr[inds, :N] = M_inv * (X_temp @ asso_leg_term_pos)
|
|
150
|
+
X_arr[inds, N:] = -M_inv * (X_temp @ asso_leg_term_neg)
|
|
129
151
|
|
|
130
|
-
# Assemble the coefficient matrix and additional terms
|
|
131
|
-
# --------------------------------------------------------------------------------------------------------------------------
|
|
132
|
-
DW = D_pos * W[None, :]
|
|
133
|
-
DW.ravel()[::N + 1] -= 1
|
|
134
|
-
alpha = M_inv[:, None] * DW
|
|
135
|
-
beta = M_inv[:, None] * D_neg * W[None, :]
|
|
136
|
-
|
|
137
|
-
# --------------------------------------------------------------------------------------------------------------------------
|
|
138
|
-
|
|
139
|
-
# Setup to solve for particular solution w.r.t. direct beam source (see section 3.6.1 of the Comprehensive Documentation)
|
|
140
|
-
# --------------------------------------------------------------------------------------------------------------------------
|
|
141
|
-
if there_is_beam_source:
|
|
142
|
-
# Generate X
|
|
143
|
-
X_temp = (
|
|
144
|
-
scalar_fac_asso_leg_term_mu0
|
|
145
|
-
* scaled_omega_l
|
|
146
|
-
* weighted_scaled_Leg_coeffs_l
|
|
147
|
-
)
|
|
148
|
-
X_pos = X_temp @ asso_leg_term_pos
|
|
149
|
-
X_neg = X_temp @ asso_leg_term_neg
|
|
150
|
-
# Signs flipped due to the extra minus sign
|
|
151
|
-
X_arr[ind, :N] = M_inv * X_pos
|
|
152
|
-
X_arr[ind, N:] = -M_inv * X_neg
|
|
153
|
-
|
|
154
|
-
# --------------------------------------------------------------------------------------------------------------------------
|
|
155
|
-
|
|
156
|
-
A_arr[ind, N:, :N] = beta
|
|
157
|
-
A_arr[ind, N:, N:] = alpha
|
|
158
|
-
no_shortcut_indices.append(ind)
|
|
159
|
-
if m_equals_0 and there_is_iso_source: # Keep the list empty if there is no isotropic source
|
|
160
|
-
no_shortcut_indices_0.append(l)
|
|
161
|
-
|
|
162
|
-
else:
|
|
163
|
-
# This is a shortcut to the diagonalization results
|
|
164
|
-
G_collect[ind, :, :] = G_trivial
|
|
165
|
-
K_collect[ind, :N] = -M_inv
|
|
166
|
-
K_collect[ind, N:] = M_inv
|
|
167
|
-
if m_equals_0 and there_is_iso_source:
|
|
168
|
-
G_inv_collect_0[l, :, :] = G_trivial
|
|
169
|
-
|
|
170
|
-
ind += 1
|
|
171
|
-
|
|
172
152
|
if len(no_shortcut_indices) > 0:
|
|
173
153
|
A_arr = A_arr[no_shortcut_indices, :, :]
|
|
174
154
|
alpha_arr = A_arr[:, N:, N:]
|
PythonicDISORT/pydisort.py
CHANGED
|
@@ -6,7 +6,7 @@ import warnings
|
|
|
6
6
|
import numpy as np
|
|
7
7
|
import scipy as sc
|
|
8
8
|
from math import pi
|
|
9
|
-
from numpy.polynomial.legendre import Legendre
|
|
9
|
+
from numpy.polynomial.legendre import Legendre, legval
|
|
10
10
|
from scipy import integrate
|
|
11
11
|
|
|
12
12
|
|
|
@@ -409,13 +409,13 @@ def pydisort(
|
|
|
409
409
|
def TMS_correction(tau, phi, is_antiderivative_wrt_tau):
|
|
410
410
|
Ntau = len(tau)
|
|
411
411
|
Nphi = len(phi)
|
|
412
|
-
# Atmospheric layer indices
|
|
413
|
-
l = np.
|
|
412
|
+
# Atmospheric layer indices: smallest l with tau <= tau_arr[l].
|
|
413
|
+
l = np.searchsorted(tau_arr, tau, side="left")
|
|
414
414
|
scaled_tau_arr_l = scaled_tau_arr_with_0[l + 1]
|
|
415
415
|
scaled_tau_arr_lm1 = scaled_tau_arr_with_0[l]
|
|
416
416
|
|
|
417
417
|
# Delta-M scaling
|
|
418
|
-
if np.any(scale_tau !=
|
|
418
|
+
if np.any(scale_tau != 1):
|
|
419
419
|
tau_dist_from_top = tau_arr[l] - tau
|
|
420
420
|
scaled_tau_dist_from_top = tau_dist_from_top * scale_tau[l]
|
|
421
421
|
scaled_tau = scaled_tau_arr_l - scaled_tau_dist_from_top
|
|
@@ -427,25 +427,15 @@ def pydisort(
|
|
|
427
427
|
nu = subroutines.atleast_2d_append(
|
|
428
428
|
subroutines.calculate_nu(mu_arr, phi, -mu0, phi0)
|
|
429
429
|
)
|
|
430
|
-
|
|
431
|
-
|
|
432
|
-
|
|
433
|
-
|
|
434
|
-
|
|
435
|
-
|
|
436
|
-
|
|
437
|
-
)
|
|
438
|
-
p_trun = np.concatenate(
|
|
439
|
-
[
|
|
440
|
-
f(nu)[:, None, :]
|
|
441
|
-
# Iterates over the 0th axis
|
|
442
|
-
for f in map(Legendre, weighted_scaled_Leg_coeffs)
|
|
443
|
-
],
|
|
444
|
-
axis=1,
|
|
445
|
-
)
|
|
430
|
+
# Vectorized Legendre evaluation: legval broadcasts the trailing-axis
|
|
431
|
+
# coefficient set, returning (NLayers, NQuad, Nphi). Move the layer axis
|
|
432
|
+
# to the middle to match the original (NQuad, NLayers, Nphi) layout.
|
|
433
|
+
# Always uses base numpy's legval (not autograd's): this matches the
|
|
434
|
+
# original `Legendre(...)(nu)` path, which was also pure numpy.
|
|
435
|
+
p_true = np.moveaxis(legval(nu, weighted_Leg_coeffs_all.T, tensor=True), 0, 1)
|
|
436
|
+
p_trun = np.moveaxis(legval(nu, weighted_scaled_Leg_coeffs.T, tensor=True), 0, 1)
|
|
446
437
|
mathscr_B = (
|
|
447
|
-
(scaled_omega_arr *
|
|
448
|
-
/ (4 * np.pi)
|
|
438
|
+
(scaled_omega_arr * I0_div_4pi)[None, :, None]
|
|
449
439
|
* (mu0 / (mu0 + mu_arr))[:, None, None]
|
|
450
440
|
* (p_true / (1 - f_arr[None, :, None]) - p_trun)
|
|
451
441
|
)
|
|
@@ -495,14 +485,16 @@ def pydisort(
|
|
|
495
485
|
any_neg = (l.max() > 0) # exists tau with layers above
|
|
496
486
|
|
|
497
487
|
# Layer boundaries/thickness in scaled-tau coordinates
|
|
498
|
-
|
|
499
|
-
tau_back = scaled_tau_arr_with_0[1:] # (NLayers,) = tau_{r+1}
|
|
488
|
+
exp_scaled_tau_arr_with_0 = prepend(np.exp(-scaled_tau_arr_with_0[1:] / mu0), NLayers, 1)
|
|
500
489
|
|
|
501
|
-
# decay[n,r] = exp(-scaled_thickness_arr[r]/mu[n])
|
|
502
|
-
|
|
490
|
+
# decay[n,r] = exp(-scaled_thickness_arr[r]/mu[n])
|
|
491
|
+
delta = scaled_thickness_arr[None, :] * M_inv[:, None]
|
|
492
|
+
decay = np.exp(-delta) # (N, NLayers)
|
|
503
493
|
|
|
504
|
-
|
|
505
|
-
|
|
494
|
+
# tau_front = scaled_tau_arr_with_0[:-1]
|
|
495
|
+
exp_taufront_mu0_inv = exp_scaled_tau_arr_with_0[:-1] # (NLayers,)
|
|
496
|
+
# tau_back = scaled_tau_arr_with_0[1:]
|
|
497
|
+
exp_tauback_mu0_inv = exp_scaled_tau_arr_with_0[1:] # (NLayers,)
|
|
506
498
|
|
|
507
499
|
if is_antiderivative_wrt_tau:
|
|
508
500
|
# (N, NLayers): multiply by this instead of dividing by (scale_tau/mu)
|
|
@@ -510,73 +502,100 @@ def pydisort(
|
|
|
510
502
|
|
|
511
503
|
# ===== POS: sum over r > l =====
|
|
512
504
|
if any_pos:
|
|
505
|
+
l_min = int(l.min()) # only Acc_pos at layers >= l_min is ever indexed
|
|
506
|
+
|
|
513
507
|
# fac_pos = 1 - exp(-scaled_thickness_arr*(1/mu + 1/mu0)) = -expm1(-scaled_thickness_arr*(...))
|
|
514
|
-
delta_pos =
|
|
515
|
-
fac_pos = -np.expm1(-delta_pos)
|
|
508
|
+
delta_pos = delta + (scaled_thickness_arr / mu0)[None, :] # (N, NLayers) >= 0
|
|
509
|
+
fac_pos = -np.expm1(-delta_pos) # (N, NLayers) in [0,1)
|
|
516
510
|
|
|
517
|
-
|
|
518
|
-
|
|
511
|
+
# NOTE: antiderivative target-layer factor `mu/scale_tau[l]` is applied
|
|
512
|
+
# OUTSIDE the accumulator at line 544, not per-source-layer inside it.
|
|
519
513
|
|
|
520
514
|
# coef_pos[n,r] multiplies exp((tau - tau_front[r])/mu - tau_front[r]/mu0) after re-referencing
|
|
521
515
|
# includes exp(-tau_front[r]/mu0)
|
|
522
|
-
coef_pos = fac_pos * exp_taufront_mu0_inv[None, :]
|
|
516
|
+
coef_pos = fac_pos * exp_taufront_mu0_inv[None, :] # (N, NLayers)
|
|
523
517
|
|
|
524
518
|
# Reverse scan to build Acc_l (referenced at tau_back[l]) INCLUDING Bpos:
|
|
525
519
|
# Acc_{r-1} = coef_pos[:,r]*Bpos[:,r,:] + decay[:,r]*Acc_r
|
|
526
|
-
|
|
527
|
-
|
|
528
|
-
|
|
529
|
-
|
|
530
|
-
|
|
531
|
-
|
|
532
|
-
|
|
533
|
-
|
|
534
|
-
|
|
535
|
-
|
|
520
|
+
# Stored array covers layers [l_min..NLayers-1]; layer r maps to index r-l_min.
|
|
521
|
+
Npos = NLayers - l_min
|
|
522
|
+
|
|
523
|
+
if autograd_compatible:
|
|
524
|
+
acc0 = np.zeros((N, Nphi))
|
|
525
|
+
acc = acc0
|
|
526
|
+
cols_rev = [acc0] # Acc_{NLayers-1} = 0
|
|
527
|
+
for r in range(NLayers - 1, l_min, -1):
|
|
528
|
+
term = coef_pos[:, r][:, None] * Bpos[:, r, :] # (N, Nphi)
|
|
529
|
+
acc = term + decay[:, r][:, None] * acc
|
|
530
|
+
cols_rev.append(acc)
|
|
531
|
+
Acc_pos = np.stack(cols_rev[::-1], axis=1) # (N, Npos, Nphi)
|
|
532
|
+
else:
|
|
533
|
+
Acc_pos = np.empty((N, Npos, Nphi))
|
|
534
|
+
Acc_pos[:, Npos - 1, :] = 0.0 # Acc_{NLayers-1}
|
|
535
|
+
for r in range(NLayers - 1, l_min, -1):
|
|
536
|
+
i_dest = r - 1 - l_min
|
|
537
|
+
i_source = r - l_min
|
|
538
|
+
np.multiply(coef_pos[:, r:r+1], Bpos[:, r, :],
|
|
539
|
+
out=Acc_pos[:, i_dest, :])
|
|
540
|
+
Acc_pos[:, i_dest, :] += decay[:, r:r+1] * Acc_pos[:, i_source, :]
|
|
536
541
|
|
|
537
542
|
# Apply to tau: multiply by exp((scaled_tau - tau_back[l]) / mu) (argument <= 0 in exact arithmetic)
|
|
538
|
-
expfac_pos = np.exp(M_inv[:, None] * (scaled_tau -
|
|
539
|
-
contrib_pos = Acc_pos[:, l, :] * expfac_pos[:, :, None]
|
|
543
|
+
expfac_pos = np.exp(M_inv[:, None] * (scaled_tau - scaled_tau_arr_with_0[1:][l])[None, :]) # (N, Ntau)
|
|
544
|
+
contrib_pos = Acc_pos[:, l - l_min, :] * expfac_pos[:, :, None] # (N, Ntau, Nphi)
|
|
545
|
+
|
|
546
|
+
if is_antiderivative_wrt_tau:
|
|
547
|
+
contrib_pos = contrib_pos * mu_div_scale_layers[:, l, None]
|
|
540
548
|
|
|
541
549
|
sol_pos = sol_pos + contrib_pos
|
|
542
550
|
|
|
543
551
|
# ===== NEG: sum over r < l =====
|
|
544
552
|
if any_neg:
|
|
553
|
+
l_max = int(l.max()) # only Acc_neg at layers <= l_max is ever indexed
|
|
554
|
+
|
|
545
555
|
# We need term_base[n,r] = exp(-tau_back[r]/mu0) - exp(-scaled_thickness_arr[r]/mu)*exp(-tau_front[r]/mu0)
|
|
546
556
|
# computed stably via expm1 using only non-positive inputs.
|
|
547
|
-
delta_neg =
|
|
548
|
-
|
|
549
|
-
em1 = np.expm1(-abs_delta) # (N, NLayers) in [-1,0], argument <= 0
|
|
557
|
+
delta_neg = delta - (scaled_thickness_arr / mu0)[None, :] # (N, NLayers), can be +/-.
|
|
558
|
+
em1 = np.expm1(-np.abs(delta_neg)) # (N, NLayers) in [-1,0], argument <= 0
|
|
550
559
|
|
|
551
560
|
exp_x1 = exp_tauback_mu0_inv[None, :] # (1, NLayers)
|
|
552
561
|
exp_x0 = decay * exp_taufront_mu0_inv[None, :] # (N, NLayers)
|
|
553
562
|
|
|
554
563
|
# if delta>=0: term = exp_x1*(1-exp(-delta)) = exp_x1*(-expm1(-delta)) = exp_x1*(-em1)
|
|
555
564
|
# if delta<0 : term = exp_x0*(exp(delta)-1) = exp_x0*expm1(delta) = exp_x0*( em1) (since delta<0 -> em1=expm1(delta))
|
|
556
|
-
term_base = np.where(delta_neg >= 0.0,
|
|
565
|
+
term_base = np.where(delta_neg >= 0.0, -em1 * exp_x1, em1 * exp_x0) # (N, NLayers)
|
|
557
566
|
|
|
558
|
-
|
|
559
|
-
|
|
560
|
-
|
|
561
|
-
else:
|
|
562
|
-
coef_neg = term_base # (N, NLayers)
|
|
567
|
+
coef_neg = term_base # (N, NLayers)
|
|
568
|
+
# NOTE: antiderivative target-layer factor `-mu/scale_tau[l]` is applied
|
|
569
|
+
# OUTSIDE the accumulator at line 600, not per-source-layer inside it.
|
|
563
570
|
|
|
564
571
|
# Forward scan to build Acc_l (referenced at tau_front[l]) INCLUDING Bneg:
|
|
565
572
|
# Acc_{r+1} = coef_neg[:,r]*Bneg[:,r,:] + decay[:,r]*Acc_r
|
|
566
|
-
|
|
567
|
-
|
|
568
|
-
|
|
569
|
-
|
|
570
|
-
|
|
571
|
-
|
|
572
|
-
|
|
573
|
-
|
|
574
|
-
|
|
575
|
-
|
|
573
|
+
# Stored array covers layers [0..l_max]; layer index unchanged.
|
|
574
|
+
Nneg = l_max + 1
|
|
575
|
+
|
|
576
|
+
if autograd_compatible:
|
|
577
|
+
acc0 = np.zeros((N, Nphi))
|
|
578
|
+
acc = acc0
|
|
579
|
+
cols = [acc0] # Acc_0 = 0
|
|
580
|
+
for r in range(0, l_max):
|
|
581
|
+
term = coef_neg[:, r][:, None] * Bneg[:, r, :] # (N, Nphi)
|
|
582
|
+
acc = term + decay[:, r][:, None] * acc
|
|
583
|
+
cols.append(acc)
|
|
584
|
+
Acc_neg = np.stack(cols, axis=1) # (N, Nneg, Nphi)
|
|
585
|
+
else:
|
|
586
|
+
Acc_neg = np.empty((N, Nneg, Nphi))
|
|
587
|
+
Acc_neg[:, 0, :] = 0.0 # Acc_0
|
|
588
|
+
for r in range(0, l_max):
|
|
589
|
+
np.multiply(coef_neg[:, r:r+1], Bneg[:, r, :],
|
|
590
|
+
out=Acc_neg[:, r + 1, :])
|
|
591
|
+
Acc_neg[:, r + 1, :] += decay[:, r:r+1] * Acc_neg[:, r, :]
|
|
576
592
|
|
|
577
593
|
# Apply to tau: multiply by exp((tau_front[l] - scaled_tau) / mu) (argument <= 0 in exact arithmetic)
|
|
578
|
-
expfac_neg = np.exp(M_inv[:, None] * (
|
|
579
|
-
contrib_neg = Acc_neg[:, l, :] * expfac_neg[:, :, None]
|
|
594
|
+
expfac_neg = np.exp(M_inv[:, None] * (scaled_tau_arr_with_0[l] - scaled_tau)[None, :]) # (N, Ntau)
|
|
595
|
+
contrib_neg = Acc_neg[:, l, :] * expfac_neg[:, :, None] # (N, Ntau, Nphi)
|
|
596
|
+
|
|
597
|
+
if is_antiderivative_wrt_tau:
|
|
598
|
+
contrib_neg = -contrib_neg * mu_div_scale_layers[:, l, None]
|
|
580
599
|
|
|
581
600
|
sol_neg = sol_neg + contrib_neg
|
|
582
601
|
|
PythonicDISORT/subroutines.py
CHANGED
|
@@ -641,60 +641,76 @@ def interpolate(u):
|
|
|
641
641
|
u_pos_interpol = sc.interpolate.BarycentricInterpolator(mu_arr_pos)
|
|
642
642
|
u_neg_interpol = sc.interpolate.BarycentricInterpolator(-mu_arr_pos)
|
|
643
643
|
|
|
644
|
+
# Value-keyed cache for the most recent (tau, phi, is_antideriv) -> u_cache.
|
|
645
|
+
# Hits when the user loops over `mu` with the same (tau, phi); otherwise re-evaluates `u`.
|
|
646
|
+
# Cache stores copies of `tau`/`phi` so in-place mutation of the user's arrays
|
|
647
|
+
# cannot return stale values (compared via `np.array_equal`).
|
|
648
|
+
# Only used on the simple-return path (skipped when Fourier_error / tau_arr are requested).
|
|
649
|
+
_cache = [None, None, None, None] # [tau_copy, phi_copy, is_antideriv, u_cache]
|
|
650
|
+
|
|
644
651
|
argcount = u.__code__.co_argcount
|
|
645
652
|
if argcount == 5: # Function is `u` instead of `u0`
|
|
646
653
|
def u_interpol(mu, tau, phi, is_antiderivative_wrt_tau=False, return_Fourier_error=False, return_tau_arr=False):
|
|
647
654
|
if not np.all(np.abs(mu) <= 1):
|
|
648
655
|
raise ValueError("mu values must be between -1 and 1.")
|
|
649
|
-
|
|
656
|
+
|
|
650
657
|
mu = np.atleast_1d(mu)
|
|
651
658
|
|
|
652
659
|
if return_Fourier_error or return_tau_arr:
|
|
653
660
|
u_outputs = u(tau, phi, is_antiderivative_wrt_tau, return_Fourier_error, return_tau_arr)
|
|
654
661
|
u_cache = u_outputs[0]
|
|
662
|
+
elif (_cache[3] is not None and is_antiderivative_wrt_tau == _cache[2]
|
|
663
|
+
and np.array_equal(tau, _cache[0]) and np.array_equal(phi, _cache[1])):
|
|
664
|
+
u_cache = _cache[3]
|
|
655
665
|
else:
|
|
656
666
|
u_cache = u(tau, phi, is_antiderivative_wrt_tau)
|
|
657
|
-
|
|
667
|
+
_cache[:] = [np.array(tau, copy=True), np.array(phi, copy=True),
|
|
668
|
+
is_antiderivative_wrt_tau, u_cache]
|
|
669
|
+
|
|
658
670
|
results = np.empty((len(mu),) + np.shape(u_cache)[1:])
|
|
659
671
|
mask_pos = mu > 0
|
|
660
672
|
mask_else = ~mask_pos
|
|
661
|
-
|
|
673
|
+
|
|
662
674
|
if np.any(mask_pos):
|
|
663
675
|
u_pos_interpol.set_yi(u_cache[:N])
|
|
664
676
|
results[mask_pos] = u_pos_interpol(mu[mask_pos])
|
|
665
677
|
if np.any(mask_else):
|
|
666
678
|
u_neg_interpol.set_yi(u_cache[N:])
|
|
667
679
|
results[mask_else] = u_neg_interpol(mu[mask_else])
|
|
668
|
-
|
|
680
|
+
|
|
669
681
|
if return_Fourier_error or return_tau_arr:
|
|
670
682
|
return (np.squeeze(results)[()],) + u_outputs[1:]
|
|
671
683
|
else:
|
|
672
684
|
return np.squeeze(results)[()]
|
|
673
|
-
|
|
685
|
+
|
|
674
686
|
elif argcount == 4: # Function is u0 instead of u
|
|
675
687
|
def u_interpol(mu, tau, is_antiderivative_wrt_tau=False, return_Fourier_error=False, return_tau_arr=False):
|
|
676
688
|
if not np.all(np.abs(mu) <= 1):
|
|
677
689
|
raise ValueError("mu values must be between -1 and 1.")
|
|
678
|
-
|
|
690
|
+
|
|
679
691
|
mu = np.atleast_1d(mu)
|
|
680
692
|
|
|
681
693
|
if return_Fourier_error or return_tau_arr:
|
|
682
694
|
u_outputs = u(tau, is_antiderivative_wrt_tau, return_Fourier_error, return_tau_arr)
|
|
683
695
|
u_cache = u_outputs[0]
|
|
696
|
+
elif (_cache[3] is not None and is_antiderivative_wrt_tau == _cache[2]
|
|
697
|
+
and np.array_equal(tau, _cache[0])):
|
|
698
|
+
u_cache = _cache[3]
|
|
684
699
|
else:
|
|
685
700
|
u_cache = u(tau, is_antiderivative_wrt_tau)
|
|
686
|
-
|
|
701
|
+
_cache[:] = [np.array(tau, copy=True), None, is_antiderivative_wrt_tau, u_cache]
|
|
702
|
+
|
|
687
703
|
results = np.empty((len(mu),) + np.shape(u_cache)[1:])
|
|
688
704
|
mask_pos = mu > 0
|
|
689
705
|
mask_else = ~mask_pos
|
|
690
|
-
|
|
706
|
+
|
|
691
707
|
if np.any(mask_pos):
|
|
692
708
|
u_pos_interpol.set_yi(u_cache[:N])
|
|
693
709
|
results[mask_pos] = u_pos_interpol(mu[mask_pos])
|
|
694
710
|
if np.any(mask_else):
|
|
695
711
|
u_neg_interpol.set_yi(u_cache[N:])
|
|
696
712
|
results[mask_else] = u_neg_interpol(mu[mask_else])
|
|
697
|
-
|
|
713
|
+
|
|
698
714
|
if return_Fourier_error or return_tau_arr:
|
|
699
715
|
return (np.squeeze(results)[()],) + u_outputs[1:]
|
|
700
716
|
else:
|
|
@@ -764,7 +780,7 @@ def _mathscr_v(tau, # Input optical depths
|
|
|
764
780
|
| Variable | Shape |
|
|
765
781
|
| ----------------------------- | ----------------------------------------- |
|
|
766
782
|
| ``tau`` | ``Ntau`` |
|
|
767
|
-
| ``scale_tau`` |
|
|
783
|
+
| ``scale_tau`` | ``NL`` |
|
|
768
784
|
| ``l`` | ``Ntau`` |
|
|
769
785
|
| ``Nscoeffs`` | scalar |
|
|
770
786
|
| ``s_poly_coeffs`` | ``NL x Nscoeffs`` |
|
|
@@ -790,11 +806,10 @@ def _mathscr_v(tau, # Input optical depths
|
|
|
790
806
|
d0 = np.einsum("lij, lj -> il", G, c0w, optimize=True)
|
|
791
807
|
|
|
792
808
|
if is_antiderivative_wrt_tau:
|
|
793
|
-
p0 = tau
|
|
809
|
+
p0 = tau
|
|
810
|
+
return (d0 / scale_tau[None, :])[:, l] * p0[None, :]
|
|
794
811
|
else:
|
|
795
|
-
|
|
796
|
-
|
|
797
|
-
return d0[:, l] * p0[None, :]
|
|
812
|
+
return d0[:, l]
|
|
798
813
|
|
|
799
814
|
# --------------------- Nscoeffs == 2 ---------------------
|
|
800
815
|
elif Nscoeffs == 2:
|
|
@@ -812,11 +827,11 @@ def _mathscr_v(tau, # Input optical depths
|
|
|
812
827
|
if is_antiderivative_wrt_tau:
|
|
813
828
|
p0 = tau * tau / 2
|
|
814
829
|
p1 = tau
|
|
815
|
-
return (d0[:, l] * p0[None, :] + d1[:, l] * p1[None, :]
|
|
830
|
+
return (d0 / scale_tau[None, :])[:, l] * p0[None, :] + (d1 / scale_tau[None, :])[:, l] * p1[None, :]
|
|
816
831
|
else:
|
|
817
832
|
p0 = tau
|
|
818
|
-
p1 = np.ones_like(tau)
|
|
819
|
-
return d0[:, l] * p0[None, :] + d1[:, l]
|
|
833
|
+
#p1 = np.ones_like(tau)
|
|
834
|
+
return d0[:, l] * p0[None, :] + d1[:, l]
|
|
820
835
|
|
|
821
836
|
# --------------------- Nscoeffs > 2 ----------------------
|
|
822
837
|
else:
|
|
@@ -855,11 +870,11 @@ def _mathscr_v(tau, # Input optical depths
|
|
|
855
870
|
powers = np.arange(Nscoeffs - 1, -1, -1)[None, :]
|
|
856
871
|
if is_antiderivative_wrt_tau:
|
|
857
872
|
p = powers + 1
|
|
858
|
-
tau_poly = (tau[:, None] ** p) /
|
|
873
|
+
tau_poly = (tau[:, None] ** p) / p
|
|
874
|
+
return np.einsum("tqp, tp -> qt", (mathscr_b / scale_tau[:, None, None])[l, :, :], tau_poly, optimize=True)
|
|
859
875
|
else:
|
|
860
876
|
tau_poly = tau[:, None] ** powers
|
|
861
|
-
|
|
862
|
-
return np.einsum("tqp, tp -> qt", mathscr_b[l, :, :], tau_poly, optimize=True)
|
|
877
|
+
return np.einsum("tqp, tp -> qt", mathscr_b[l, :, :], tau_poly, optimize=True)
|
|
863
878
|
|
|
864
879
|
|
|
865
880
|
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.1
|
|
2
2
|
Name: PythonicDISORT
|
|
3
|
-
Version: 1.
|
|
3
|
+
Version: 1.7
|
|
4
4
|
Summary: Discrete Ordinates Solver for the (1D) Radiative Transfer Equation in a single or multi-layer atmosphere.
|
|
5
5
|
Author-email: Dion HO Jia Xu <dh3065@columbia.edu>
|
|
6
6
|
Project-URL: Homepage, https://github.com/LDEO-CREW/Pythonic-DISORT
|
|
@@ -0,0 +1,11 @@
|
|
|
1
|
+
PythonicDISORT/__init__.py,sha256=V1obNzaFgqInxqYJ5h_r6JHbZRxx4VM5VgzCIS6EEzI,84
|
|
2
|
+
PythonicDISORT/_assemble_intensity_and_fluxes.py,sha256=-IrRbIxW-AeEAfLpIK-J1jQI3rtK4MjNm4BxhCgXzf0,29774
|
|
3
|
+
PythonicDISORT/_solve_for_coeffs.py,sha256=hzg8-II3b-vaB6U_3zxMfeV59FCTf52MEN3yknIqW18,20959
|
|
4
|
+
PythonicDISORT/_solve_for_gen_and_part_sols.py,sha256=3OFDQ7eT_xSeI4TUJ0HBNdJIYuJR-JWOA5P5pjfBjik,13978
|
|
5
|
+
PythonicDISORT/pydisort.py,sha256=fImCsduUL15xVRus_jfBaW1LD3SUTRPFyk1o1L3tSmA,39563
|
|
6
|
+
PythonicDISORT/subroutines.py,sha256=tVydhXBsqXpf-FDlVw90RHpxKQY_V3E6ziAobuGJcXw,36378
|
|
7
|
+
pythonicdisort-1.7.dist-info/LICENSE.md,sha256=q1i1QVh0OLzK0fj6kaNRNWsM_KZca66JwWgWwbuCPlA,1092
|
|
8
|
+
pythonicdisort-1.7.dist-info/METADATA,sha256=KYQtE9inxj6NRsRQGykYIrBh-dp4dVIdWZqLfgiqssA,8699
|
|
9
|
+
pythonicdisort-1.7.dist-info/WHEEL,sha256=BNRMDyzLkkcmlv0J8ppDQkk2VED33SesJDynr9ED1gc,91
|
|
10
|
+
pythonicdisort-1.7.dist-info/top_level.txt,sha256=aX6kCeI1E7naZz3xfw4_mIbWh5KBh9dmF3FIvgP5m3Y,15
|
|
11
|
+
pythonicdisort-1.7.dist-info/RECORD,,
|
|
@@ -1,11 +0,0 @@
|
|
|
1
|
-
PythonicDISORT/__init__.py,sha256=V1obNzaFgqInxqYJ5h_r6JHbZRxx4VM5VgzCIS6EEzI,84
|
|
2
|
-
PythonicDISORT/_assemble_intensity_and_fluxes.py,sha256=nN2Kn7UCB6tYlVjKj1Hk_IwNO420uibsSD6o9Lx9X5o,29520
|
|
3
|
-
PythonicDISORT/_solve_for_coeffs.py,sha256=5hXhKi_wkEBAcRmoeauGWDygmH3aB0g9QPZN7I2gaIM,19529
|
|
4
|
-
PythonicDISORT/_solve_for_gen_and_part_sols.py,sha256=5p5AlosW2ALjUBn6aEcPj_1pDpfxMn4MwwGorf9WLUY,14177
|
|
5
|
-
PythonicDISORT/pydisort.py,sha256=GRVU2cEzM1ySuPtSEeYtmRdEPnLa3vLzfA8xd1vb95w,37860
|
|
6
|
-
PythonicDISORT/subroutines.py,sha256=YscXrzzT2pKexe4S-30dfpMX69oUqyvQPnaMu8C1uGc,35223
|
|
7
|
-
pythonicdisort-1.6.dist-info/LICENSE.md,sha256=q1i1QVh0OLzK0fj6kaNRNWsM_KZca66JwWgWwbuCPlA,1092
|
|
8
|
-
pythonicdisort-1.6.dist-info/METADATA,sha256=uc2na6BfMyUgBFE1GSQe3RmekJa4EqIV2t-EK8ZPOAk,8699
|
|
9
|
-
pythonicdisort-1.6.dist-info/WHEEL,sha256=WnJ8fYhv8N4SYVK2lLYNI6N0kVATA7b0piVUNvqIIJE,91
|
|
10
|
-
pythonicdisort-1.6.dist-info/top_level.txt,sha256=aX6kCeI1E7naZz3xfw4_mIbWh5KBh9dmF3FIvgP5m3Y,15
|
|
11
|
-
pythonicdisort-1.6.dist-info/RECORD,,
|
|
File without changes
|
|
File without changes
|