lightweaver 0.16.1__cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.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 (69) hide show
  1. lightweaver/Data/AbundancesAsplund09.pickle +0 -0
  2. lightweaver/Data/AtomicMassesNames.pickle +0 -0
  3. lightweaver/Data/Barklem_dfdata.dat +41 -0
  4. lightweaver/Data/Barklem_pddata.dat +40 -0
  5. lightweaver/Data/Barklem_spdata.dat +46 -0
  6. lightweaver/Data/DefaultMolecules/C2.molecule +27 -0
  7. lightweaver/Data/DefaultMolecules/CH/CH_X-A.asc +46409 -0
  8. lightweaver/Data/DefaultMolecules/CH/CH_X-A_12.asc +28322 -0
  9. lightweaver/Data/DefaultMolecules/CH/CH_X-B.asc +4272 -0
  10. lightweaver/Data/DefaultMolecules/CH/CH_X-B_12.asc +2583 -0
  11. lightweaver/Data/DefaultMolecules/CH/CH_X-C.asc +20916 -0
  12. lightweaver/Data/DefaultMolecules/CH/CH_X-C_12.asc +13106 -0
  13. lightweaver/Data/DefaultMolecules/CH.molecule +35 -0
  14. lightweaver/Data/DefaultMolecules/CN.molecule +30 -0
  15. lightweaver/Data/DefaultMolecules/CO/vmax=3_Jmax=49_dv=1_26 +296 -0
  16. lightweaver/Data/DefaultMolecules/CO/vmax=9_Jmax=120_dv=1_26 +2162 -0
  17. lightweaver/Data/DefaultMolecules/CO.molecule +30 -0
  18. lightweaver/Data/DefaultMolecules/CO_NLTE.molecule +29 -0
  19. lightweaver/Data/DefaultMolecules/CaH.molecule +29 -0
  20. lightweaver/Data/DefaultMolecules/H2+.molecule +27 -0
  21. lightweaver/Data/DefaultMolecules/H2.molecule +27 -0
  22. lightweaver/Data/DefaultMolecules/H2O.molecule +27 -0
  23. lightweaver/Data/DefaultMolecules/HF.molecule +29 -0
  24. lightweaver/Data/DefaultMolecules/LiH.molecule +27 -0
  25. lightweaver/Data/DefaultMolecules/MgH.molecule +34 -0
  26. lightweaver/Data/DefaultMolecules/N2.molecule +28 -0
  27. lightweaver/Data/DefaultMolecules/NH.molecule +27 -0
  28. lightweaver/Data/DefaultMolecules/NO.molecule +27 -0
  29. lightweaver/Data/DefaultMolecules/O2.molecule +27 -0
  30. lightweaver/Data/DefaultMolecules/OH.molecule +27 -0
  31. lightweaver/Data/DefaultMolecules/SiO.molecule +26 -0
  32. lightweaver/Data/DefaultMolecules/TiO.molecule +30 -0
  33. lightweaver/Data/Quadratures.pickle +0 -0
  34. lightweaver/Data/pf_Kurucz.input +0 -0
  35. lightweaver/DefaultIterSchemes/.placeholder +0 -0
  36. lightweaver/DefaultIterSchemes/SimdImpl_AVX2FMA.cpython-312-x86_64-linux-gnu.so +0 -0
  37. lightweaver/DefaultIterSchemes/SimdImpl_AVX512.cpython-312-x86_64-linux-gnu.so +0 -0
  38. lightweaver/DefaultIterSchemes/SimdImpl_SSE2.cpython-312-x86_64-linux-gnu.so +0 -0
  39. lightweaver/LwCompiled.cpython-312-x86_64-linux-gnu.so +0 -0
  40. lightweaver/__init__.py +33 -0
  41. lightweaver/atmosphere.py +1767 -0
  42. lightweaver/atomic_model.py +852 -0
  43. lightweaver/atomic_set.py +1286 -0
  44. lightweaver/atomic_table.py +653 -0
  45. lightweaver/barklem.py +151 -0
  46. lightweaver/benchmark.py +113 -0
  47. lightweaver/broadening.py +605 -0
  48. lightweaver/collisional_rates.py +337 -0
  49. lightweaver/config.py +106 -0
  50. lightweaver/constants.py +22 -0
  51. lightweaver/crtaf.py +197 -0
  52. lightweaver/fal.py +440 -0
  53. lightweaver/iterate_ctx.py +241 -0
  54. lightweaver/iteration_update.py +134 -0
  55. lightweaver/libenkiTS.so +0 -0
  56. lightweaver/molecule.py +225 -0
  57. lightweaver/multi.py +113 -0
  58. lightweaver/nr_update.py +106 -0
  59. lightweaver/rh_atoms.py +19743 -0
  60. lightweaver/simd_management.py +42 -0
  61. lightweaver/utils.py +509 -0
  62. lightweaver/version.py +34 -0
  63. lightweaver/wittmann.py +1375 -0
  64. lightweaver/zeeman.py +157 -0
  65. lightweaver-0.16.1.dist-info/METADATA +81 -0
  66. lightweaver-0.16.1.dist-info/RECORD +69 -0
  67. lightweaver-0.16.1.dist-info/WHEEL +6 -0
  68. lightweaver-0.16.1.dist-info/licenses/LICENSE +21 -0
  69. lightweaver-0.16.1.dist-info/top_level.txt +1 -0
@@ -0,0 +1,1767 @@
1
+ import numbers
2
+ import pickle
3
+ from copy import copy
4
+ from dataclasses import dataclass
5
+ from enum import Enum, auto
6
+ from typing import TYPE_CHECKING, Optional, Sequence, Union, cast
7
+
8
+ import astropy.units as u
9
+ import numpy as np
10
+ from numpy.polynomial.legendre import leggauss
11
+
12
+ import lightweaver.constants as Const
13
+ from .atomic_table import (AtomicAbundance, DefaultAtomicAbundance,
14
+ PeriodicTable)
15
+ from .utils import (ConvergenceError, check_shape_exception, get_data_path,
16
+ view_flatten)
17
+ from .wittmann import Wittmann, cgs
18
+
19
+ if TYPE_CHECKING:
20
+ from .LwCompiled import LwSpectrum
21
+
22
+ class ScaleType(Enum):
23
+ '''
24
+ Atmospheric scales used in the definition of 1D atmospheres to allow the
25
+ correct conversion to a height based system.
26
+ Options:
27
+
28
+ - `Geometric`
29
+
30
+ - `ColumnMass`
31
+
32
+ - `Tau500`
33
+
34
+ '''
35
+ Geometric = 0
36
+ ColumnMass = auto()
37
+ Tau500 = auto()
38
+
39
+ class BoundaryCondition:
40
+ '''
41
+ Base class for boundary conditions.
42
+
43
+ Defines the interface; do not use directly.
44
+
45
+ Attributes
46
+ ----------
47
+ These attributes are only available after set_required_angles has been called.
48
+
49
+ mux : np.ndarray
50
+ The mu_x to return from compute_bc (in order).
51
+ muy : np.ndarray
52
+ The mu_y to return from compute_bc (in order).
53
+ muz : np.ndarray
54
+ The mu_z to return from compute_bc (in order).
55
+ indexVector : np.ndarray
56
+ A 2D array of integer shape (mu, toObs) - where mu is the mu index on
57
+ the associated atmosphere - relating each index of the second (Nrays)
58
+ axis of a pair of (mu, toObs). Used to construct and destructure this
59
+ array.
60
+
61
+ '''
62
+ def compute_bc(self, atmos: 'Atmosphere', spect: 'LwSpectrum') -> np.ndarray:
63
+ '''
64
+ Called when the radiation boundary condition is needed by the backend.
65
+
66
+ Parameters
67
+ ----------
68
+ atmos : Atmosphere
69
+ The atmospheric object in which to compute the radiation.
70
+ spect : LwSpectrum
71
+ The computational spectrum object provided by the Context.
72
+
73
+ Returns
74
+ -------
75
+ result : np.ndarray
76
+ This function needs to return a contiguous array of shape [Nwave,
77
+ Nrays, Nbc], where Nwave is the number of wavelengths in the
78
+ wavelength grid, Nrays is the number of rays in the angular
79
+ quadrature (also including up/down directions) ordered as
80
+ specified by the mux/y/z and indexVector variables on this
81
+ object, Nbc is the number of spatial positions the boundary
82
+ condition needs to be defined at ordered in a flattened [Nz, Ny,
83
+ Nx] fashion. (dtype: <f8)
84
+
85
+ '''
86
+ raise NotImplementedError
87
+
88
+ def set_required_angles(self, mux, muy, muz, indexVector):
89
+ '''
90
+ The angles (and their ordering) to be used for this boundary
91
+ condition (in the case of a callable)
92
+ '''
93
+ self.mux = mux
94
+ self.muy = muy
95
+ self.muz = muz
96
+ self.indexVector = indexVector
97
+
98
+ class NoBc(BoundaryCondition):
99
+ '''
100
+ Indicates no boundary condition on the axis because it is invalid for the
101
+ current simulation.
102
+ Used only by the backend.
103
+ '''
104
+ pass
105
+
106
+ class ZeroRadiation(BoundaryCondition):
107
+ '''
108
+ Zero radiation boundary condition.
109
+ Commonly used for coronal situations.
110
+ '''
111
+ pass
112
+
113
+ class ThermalisedRadiation(BoundaryCondition):
114
+ '''
115
+ Thermalised radiation (blackbody) boundary condition.
116
+ Commonly used for photospheric situations.
117
+ '''
118
+ pass
119
+
120
+ class PeriodicRadiation(BoundaryCondition):
121
+ '''
122
+ Periodic boundary condition.
123
+ Commonly used on the x-axis in 2D simulations.
124
+ '''
125
+ pass
126
+
127
+ def get_top_pressure(eos: Wittmann, temp, ne=None, rho=None):
128
+ '''
129
+ Return a pressure for the top of atmosphere.
130
+ For internal use.
131
+
132
+ In order this is deduced from:
133
+ - the electron density `ne` [m-3], if provided
134
+ - the mass density `rho` [kg m-3], if provided
135
+ - the electron pressure present in FALC
136
+
137
+ Returns
138
+ -------
139
+ pressure : float
140
+ pressure IN CGS [dyn cm-2]
141
+
142
+ '''
143
+ if ne is not None:
144
+ pe = (ne << u.Unit('m-3')).to('cm-3').value * cgs.BK * temp
145
+ return eos.pg_from_pe(temp, pe)
146
+ elif rho is not None:
147
+ return eos.pg_from_rho(temp, (rho << u.Unit('kg m-3')).to('g cm-3').value)
148
+
149
+ pgasCgs = np.array([0.70575286, 0.59018545, 0.51286639, 0.43719268, 0.37731009,
150
+ 0.33516886, 0.31342915, 0.30604891, 0.30059491, 0.29207645,
151
+ 0.2859011 , 0.28119224, 0.27893046, 0.27949676, 0.28299726,
152
+ 0.28644693, 0.28825946, 0.29061192, 0.29340255, 0.29563072,
153
+ 0.29864548, 0.30776456, 0.31825915, 0.32137574, 0.3239401 ,
154
+ 0.32622212, 0.32792196, 0.3292243 , 0.33025437, 0.33146736,
155
+ 0.3319676 , 0.33217821, 0.3322355 , 0.33217166, 0.33210297,
156
+ 0.33203833, 0.33198508])
157
+ tempCoord = np.array([ 7600., 7780., 7970., 8273., 8635., 8988., 9228.,
158
+ 9358., 9458., 9587., 9735., 9983., 10340., 10850.,
159
+ 11440., 12190., 13080., 14520., 16280., 17930., 20420.,
160
+ 24060., 27970., 32150., 36590., 41180., 45420., 49390.,
161
+ 53280., 60170., 66150., 71340., 75930., 83890., 90820.,
162
+ 95600., 100000.])
163
+
164
+ ptop = np.interp(temp, tempCoord, pgasCgs)
165
+ return ptop
166
+
167
+ @dataclass
168
+ class Stratifications:
169
+ '''
170
+ Stores the optional derived z-stratifications of an atmospheric model.
171
+
172
+ Attributes
173
+ ----------
174
+ cmass : np.ndarray
175
+ Column mass [kg m-2].
176
+ tauRef : np.ndarray
177
+ Reference optical depth at 500 nm.
178
+ '''
179
+ cmass: np.ndarray
180
+ tauRef: np.ndarray
181
+
182
+ def dimensioned_view(self, shape) -> 'Stratifications':
183
+ '''
184
+ Makes an instance of `Stratifications` reshaped to the provided
185
+ shape for multi-dimensional atmospheres.
186
+ For internal use.
187
+
188
+ Parameters
189
+ ----------
190
+ shape : tuple
191
+ Shape to reform the stratifications, provided by
192
+ `Layout.dimensioned_shape`.
193
+
194
+ Returns
195
+ -------
196
+ stratifications : Stratifications
197
+ Reshaped stratifications.
198
+ '''
199
+ strat = copy(self)
200
+ strat.cmass = self.cmass.reshape(shape)
201
+ strat.tauRef = self.tauRef.reshape(shape)
202
+ return strat
203
+
204
+ def unit_view(self) -> 'Stratifications':
205
+ '''
206
+ Makes an instance of `Stratifications` with the correct `astropy.units`
207
+ For internal use.
208
+
209
+ Returns
210
+ -------
211
+ stratifications : Stratifications
212
+ The same data with units applied.
213
+ '''
214
+ strat = copy(self)
215
+ strat.cmass = self.cmass << u.kg / u.m**2
216
+ strat.tauRef = self.tauRef << u.dimensionless_unscaled
217
+ return strat
218
+
219
+ def dimensioned_unit_view(self, shape) -> 'Stratifications':
220
+ '''
221
+ Makes an instance of `Stratifications` reshaped to the provided shape
222
+ with the correct `astropy.units` for multi-dimensional atmospheres.
223
+ For internal use.
224
+
225
+ Parameters
226
+ ----------
227
+ shape : tuple
228
+ Shape to reform the stratifications, provided by
229
+ `Layout.dimensioned_shape`.
230
+
231
+ Returns
232
+ -------
233
+ stratifications : Stratifications
234
+ Reshaped stratifications with units.
235
+ '''
236
+ strat = self.dimensioned_view(shape)
237
+ return strat.unit_view()
238
+
239
+ @dataclass
240
+ class Layout:
241
+ '''
242
+ Storage for basic atmospheric parameters whose presence is determined by
243
+ problem dimensionality, boundary conditions and optional stratifications.
244
+
245
+ Attributes
246
+ ----------
247
+ Ndim : int
248
+ Number of dimensions in model.
249
+
250
+ x : np.ndarray
251
+ Ordinates of grid points along the x-axis (present for Ndim >= 2) [m].
252
+ y : np.ndarray
253
+ Ordinates of grid points along the y-axis (present for Ndim == 3) [m].
254
+ z : np.ndarray
255
+ Ordinates of grid points along the z-axis (present for all Ndim) [m].
256
+ vx : np.ndarray
257
+ x component of plasma velocity (optional for Ndim < 2) [m/s].
258
+ vy : np.ndarray
259
+ y component of plasma velocity (optional for Ndim < 3) [m/s].
260
+ vz : np.ndarray
261
+ z component of plasma velocity (present for all Ndim) [m/s]. Aliased to
262
+ `vlos` when `Ndim==1`
263
+ xLowerBc : BoundaryCondition
264
+ Boundary condition for the plane of minimal x-coordinate.
265
+ xUpperBc : BoundaryCondition
266
+ Boundary condition for the plane of maximal x-coordinate.
267
+ yLowerBc : BoundaryCondition
268
+ Boundary condition for the plane of minimal y-coordinate.
269
+ yUpperBc : BoundaryCondition
270
+ Boundary condition for the plane of maximal y-coordinate.
271
+ zLowerBc : BoundaryCondition
272
+ Boundary condition for the plane of minimal z-coordinate.
273
+ zUpperBc : BoundaryCondition
274
+ Boundary condition for the plane of maximal z-coordinate.
275
+ '''
276
+
277
+ Ndim: int
278
+ x: np.ndarray
279
+ y: np.ndarray
280
+ z: np.ndarray
281
+ vx: np.ndarray
282
+ vy: np.ndarray
283
+ vz: np.ndarray
284
+ xLowerBc: BoundaryCondition
285
+ xUpperBc: BoundaryCondition
286
+ yLowerBc: BoundaryCondition
287
+ yUpperBc: BoundaryCondition
288
+ zLowerBc: BoundaryCondition
289
+ zUpperBc: BoundaryCondition
290
+ stratifications: Optional[Stratifications] = None
291
+
292
+ @classmethod
293
+ def make_1d(
294
+ cls,
295
+ z: np.ndarray,
296
+ vz: np.ndarray,
297
+ lowerBc: BoundaryCondition,
298
+ upperBc: BoundaryCondition,
299
+ stratifications: Optional[Stratifications]=None,
300
+ vx: Optional[np.ndarray]=None,
301
+ vy: Optional[np.ndarray]=None,
302
+ ) -> 'Layout':
303
+ '''
304
+ Construct 1D Layout.
305
+ '''
306
+
307
+ if vx is None:
308
+ vx = np.array(())
309
+ if vy is None:
310
+ vy = np.array(())
311
+
312
+ return cls(Ndim=1, x=np.array(()), y=np.array(()),
313
+ z=z, vx=vx, vy=vy,
314
+ vz=vz, xLowerBc=NoBc(), xUpperBc=NoBc(),
315
+ yLowerBc=NoBc(), yUpperBc=NoBc(),
316
+ zLowerBc=lowerBc, zUpperBc=upperBc,
317
+ stratifications=stratifications)
318
+
319
+ @classmethod
320
+ def make_2d(
321
+ cls,
322
+ x: np.ndarray,
323
+ z: np.ndarray,
324
+ vx: np.ndarray,
325
+ vz: np.ndarray,
326
+ xLowerBc: BoundaryCondition,
327
+ xUpperBc: BoundaryCondition,
328
+ zLowerBc: BoundaryCondition,
329
+ zUpperBc: BoundaryCondition,
330
+ stratifications: Optional[Stratifications]=None,
331
+ vy: Optional[np.ndarray]=None
332
+ ) -> 'Layout':
333
+ '''
334
+ Construct 2D Layout.
335
+ '''
336
+ if vy is None:
337
+ vy = np.array(())
338
+
339
+ Bc = BoundaryCondition
340
+ return cls(Ndim=2, x=x, y=np.array(()), z=z,
341
+ vx=vx, vy=vy, vz=vz,
342
+ xLowerBc=xLowerBc, xUpperBc=xUpperBc,
343
+ yLowerBc=NoBc(), yUpperBc=NoBc(),
344
+ zLowerBc=zLowerBc, zUpperBc=zUpperBc,
345
+ stratifications=stratifications)
346
+
347
+ @classmethod
348
+ def make_3d(cls, x: np.ndarray, y: np.ndarray, z: np.ndarray,
349
+ vx: np.ndarray, vy: np.ndarray, vz: np.ndarray,
350
+ xLowerBc: BoundaryCondition, xUpperBc: BoundaryCondition,
351
+ yLowerBc: BoundaryCondition, yUpperBc: BoundaryCondition,
352
+ zLowerBc: BoundaryCondition, zUpperBc: BoundaryCondition,
353
+ stratifications: Optional[Stratifications]=None) -> 'Layout':
354
+ '''
355
+ Construct 3D Layout.
356
+ '''
357
+
358
+ return cls(Ndim=3, x=x, y=y, z=z,
359
+ vx=vx, vy=vy, vz=vz,
360
+ xLowerBc=xLowerBc, xUpperBc=xUpperBc,
361
+ yLowerBc=yLowerBc, yUpperBc=yUpperBc,
362
+ zLowerBc=zLowerBc, zUpperBc=zUpperBc,
363
+ stratifications=stratifications)
364
+
365
+ @property
366
+ def Nx(self) -> int:
367
+ '''
368
+ Number of grid points along the x-axis.
369
+ '''
370
+ return self.x.shape[0]
371
+
372
+ @property
373
+ def Ny(self) -> int:
374
+ '''
375
+ Number of grid points along the y-axis.
376
+ '''
377
+ return self.y.shape[0]
378
+
379
+ @property
380
+ def Nz(self) -> int:
381
+ '''
382
+ Number of grid points along the z-axis.
383
+ '''
384
+ return self.z.shape[0]
385
+
386
+ @property
387
+ def Noutgoing(self) -> int:
388
+ '''
389
+ Number of grid points at which the outgoing radiation is computed.
390
+ '''
391
+ return max(1, self.Nx, self.Nx * self.Ny)
392
+
393
+ @property
394
+ def vlos(self) -> np.ndarray:
395
+ if self.Ndim > 1:
396
+ raise ValueError('vlos is ambiguous when Ndim > 1, use vx, vy, or vz instead.')
397
+ return self.vz
398
+
399
+ @property
400
+ def Nspace(self) -> int:
401
+ '''
402
+ Number of spatial points present in the grid.
403
+ '''
404
+ if self.Ndim == 1:
405
+ return self.Nz
406
+ elif self.Ndim == 2:
407
+ return self.Nx * self.Nz
408
+ elif self.Ndim == 3:
409
+ return self.Nx * self.Ny * self.Nz
410
+ else:
411
+ raise ValueError('Invalid Ndim: %d, check geometry initialisation' % self.Ndim)
412
+
413
+ @property
414
+ def tauRef(self):
415
+ '''
416
+ Alias to `self.stratifications.tauRef`, if computed.
417
+ '''
418
+ if self.stratifications is not None:
419
+ return self.stratifications.tauRef
420
+ else:
421
+ raise ValueError('tauRef not computed for this Atmosphere')
422
+
423
+ @property
424
+ def cmass(self):
425
+ '''
426
+ Alias to `self.stratifications.cmass`, if computed.
427
+ '''
428
+ if self.stratifications is not None:
429
+ return self.stratifications.cmass
430
+ else:
431
+ raise ValueError('tauRef not computed for this Atmosphere')
432
+
433
+ @property
434
+ def dimensioned_shape(self):
435
+ '''
436
+ Tuple defining the shape to which the arrays of atmospheric paramters
437
+ can be reshaped to be indexed in a 1/2/3D fashion.
438
+ '''
439
+ if self.Ndim == 1:
440
+ shape = (self.Nz,)
441
+ elif self.Ndim == 2:
442
+ shape = (self.Nz, self.Nx)
443
+ elif self.Ndim == 3:
444
+ shape = (self.Nz, self.Ny, self.Nx)
445
+ else:
446
+ raise ValueError('Unreasonable Ndim (%d)' % self.Ndim)
447
+ return shape
448
+
449
+ def dimensioned_view(self) -> 'Layout':
450
+ '''
451
+ Returns a view over the contents of Layout reshaped so all data has
452
+ the correct (1/2/3D) dimensionality for the atmospheric model, as
453
+ these are all stored under a flat scheme.
454
+ '''
455
+ layout = copy(self)
456
+ shape = self.dimensioned_shape
457
+ if self.stratifications is not None:
458
+ layout.stratifications = self.stratifications.dimensioned_view(shape)
459
+ if self.vx.size > 0:
460
+ layout.vx = self.vx.reshape(shape)
461
+ if self.vy.size > 0:
462
+ layout.vy = self.vy.reshape(shape)
463
+ if self.vz.size > 0:
464
+ layout.vz = self.vz.reshape(shape)
465
+ return layout
466
+
467
+ def unit_view(self) -> 'Layout':
468
+ '''
469
+ Returns a view over the contents of the Layout with the correct
470
+ `astropy.units`.
471
+ '''
472
+ layout = copy(self)
473
+ layout.x = self.x << u.m
474
+ layout.y = self.y << u.m
475
+ layout.z = self.z << u.m
476
+ layout.vx = self.vx << u.m / u.s
477
+ layout.vy = self.vy << u.m / u.s
478
+ layout.vz = self.vz << u.m / u.s
479
+ if self.stratifications is not None:
480
+ layout.stratifications = self.stratifications.unit_view()
481
+ return layout
482
+
483
+ def dimensioned_unit_view(self) -> 'Layout':
484
+ '''
485
+ Returns a view over the contents of Layout reshaped so all data has
486
+ the correct (1/2/3D) dimensionality for the atmospheric model, and
487
+ the correct `astropy.units`.
488
+ '''
489
+ layout = self.dimensioned_view()
490
+ return layout.unit_view()
491
+
492
+ @dataclass
493
+ class Atmosphere:
494
+ '''
495
+ Storage for all atmospheric data. These arrays will be shared directly
496
+ with the backend, so a modification here also modifies the data seen by
497
+ the backend. Be careful to modify these arrays *in place*, as their data
498
+ is shared by direct memory reference. Use the class methods to construct
499
+ atmospheres of different dimensionality.
500
+
501
+ Attributes
502
+ ----------
503
+ structure : Layout
504
+ A layout structure holding the atmospheric stratification, and
505
+ velocity description.
506
+ temperature : np.ndarray
507
+ The atmospheric temperature structure.
508
+ vturb : np.ndarray
509
+ The atmospheric microturbulent velocity structure.
510
+ ne : np.ndarray
511
+ The electron density structure in the atmosphere.
512
+ nHTot : np.ndarray
513
+ The total hydrogen number density distribution throughout the
514
+ atmosphere.
515
+ B : np.ndarray, optional
516
+ The magnitude of the stratified magnetic field throughout the
517
+ atmosphere (Tesla).
518
+ gammaB : np.ndarray, optional
519
+ Co-altitude (latitude) of magnetic field vector (radians) throughout the
520
+ atmosphere from the local vertical.
521
+ chiB : np.ndarray, optional
522
+ Azimuth of magnetic field vector (radians) in the x-y plane, measured
523
+ from the x-axis.
524
+ '''
525
+
526
+ structure: Layout
527
+ temperature: np.ndarray
528
+ vturb: np.ndarray
529
+ ne: np.ndarray
530
+ nHTot: np.ndarray
531
+ B: Optional[np.ndarray] = None
532
+ gammaB: Optional[np.ndarray] = None
533
+ chiB: Optional[np.ndarray] = None
534
+
535
+ @property
536
+ def Ndim(self) -> int:
537
+ '''
538
+ Ndim : int
539
+ The dimensionality (1, 2, or 3) of the atmospheric model.
540
+ '''
541
+ return self.structure.Ndim
542
+
543
+ @property
544
+ def Nx(self) -> int:
545
+ '''
546
+ Nx : int
547
+ The number of points in the x-direction discretisation.
548
+ '''
549
+ return self.structure.Nx
550
+
551
+ @property
552
+ def Ny(self) -> int:
553
+ '''
554
+ Ny : int
555
+ The number of points in the y-direction discretisation.
556
+ '''
557
+ return self.structure.Ny
558
+
559
+ @property
560
+ def Nz(self) -> int:
561
+ '''
562
+ Nz : int
563
+ The number of points in the y-direction discretisation.
564
+ '''
565
+ return self.structure.Nz
566
+
567
+ @property
568
+ def Noutgoing(self) -> int:
569
+ '''
570
+ Noutgoing : int
571
+ The number of cells at the top of the atmosphere (that each produce a
572
+ spectrum).
573
+ '''
574
+ return self.structure.Noutgoing
575
+
576
+ @property
577
+ def vx(self) -> np.ndarray:
578
+ '''
579
+ vx : np.ndarray
580
+ x component of plasma velocity (present for Ndim >= 2) [m/s].
581
+ '''
582
+ return self.structure.vx
583
+
584
+ @property
585
+ def vy(self) -> np.ndarray:
586
+ '''
587
+ vy : np.ndarray
588
+ y component of plasma velocity (present for Ndim == 3) [m/s].
589
+ '''
590
+ return self.structure.vy
591
+
592
+ @property
593
+ def vz(self) -> np.ndarray:
594
+ '''
595
+ vz : np.ndarray
596
+ z component of plasma velocity (present for all Ndim) [m/s]. Aliased
597
+ to `vlos` when `Ndim==1`
598
+ '''
599
+ return self.structure.vz
600
+
601
+ @property
602
+ def vlos(self) -> np.ndarray:
603
+ '''
604
+ vz : np.ndarray
605
+ z component of plasma velocity (present for all Ndim) [m/s]. Only
606
+ available when Ndim==1`.
607
+ '''
608
+ return self.structure.vlos
609
+
610
+ @property
611
+ def cmass(self) -> np.ndarray:
612
+ '''
613
+ cmass : np.ndarray
614
+ Column mass [kg m-2].
615
+ '''
616
+ return self.structure.cmass
617
+
618
+ @property
619
+ def tauRef(self) -> np.ndarray:
620
+ '''
621
+ tauRef : np.ndarray
622
+ Reference optical depth at 500 nm.
623
+ '''
624
+ return self.structure.tauRef
625
+
626
+ @property
627
+ def height(self) -> np.ndarray:
628
+ return self.structure.z
629
+
630
+ @property
631
+ def x(self) -> np.ndarray:
632
+ '''
633
+ x : np.ndarray
634
+ Ordinates of grid points along the x-axis (present for Ndim >= 2) [m].
635
+ '''
636
+ return self.structure.x
637
+
638
+ @property
639
+ def y(self) -> np.ndarray:
640
+ '''
641
+ y : np.ndarray
642
+ Ordinates of grid points along the y-axis (present for Ndim == 3) [m].
643
+ '''
644
+ return self.structure.y
645
+
646
+ @property
647
+ def z(self) -> np.ndarray:
648
+ '''
649
+ z : np.ndarray
650
+ Ordinates of grid points along the z-axis (present for all Ndim) [m].
651
+ '''
652
+ return self.structure.z
653
+
654
+ @property
655
+ def zLowerBc(self) -> BoundaryCondition:
656
+ '''
657
+ zLowerBc : BoundaryCondition
658
+ Boundary condition for the plane of minimal z-coordinate.
659
+ '''
660
+ return self.structure.zLowerBc
661
+
662
+ @property
663
+ def zUpperBc(self) -> BoundaryCondition:
664
+ '''
665
+ zUpperBc : BoundaryCondition
666
+ Boundary condition for the plane of maximal z-coordinate.
667
+ '''
668
+ return self.structure.zUpperBc
669
+
670
+ @property
671
+ def yLowerBc(self) -> BoundaryCondition:
672
+ '''
673
+ yLowerBc : BoundaryCondition
674
+ Boundary condition for the plane of minimal y-coordinate.
675
+ '''
676
+ return self.structure.yLowerBc
677
+
678
+ @property
679
+ def yUpperBc(self) -> BoundaryCondition:
680
+ '''
681
+ yUpperBc : BoundaryCondition
682
+ Boundary condition for the plane of maximal y-coordinate.
683
+ '''
684
+ return self.structure.yUpperBc
685
+
686
+ @property
687
+ def xLowerBc(self) -> BoundaryCondition:
688
+ '''
689
+ xLowerBc : BoundaryCondition
690
+ Boundary condition for the plane of minimal x-coordinate.
691
+ '''
692
+ return self.structure.xLowerBc
693
+
694
+ @property
695
+ def xUpperBc(self) -> BoundaryCondition:
696
+ '''
697
+ xUpperBc : BoundaryCondition
698
+ Boundary condition for the plane of maximal x-coordinate.
699
+ '''
700
+ return self.structure.xUpperBc
701
+
702
+ @property
703
+ def Nspace(self):
704
+ '''
705
+ Nspace : int
706
+ Total number of points in the atmospheric spatial discretistaion.
707
+ '''
708
+ return self.structure.Nspace
709
+
710
+ @property
711
+ def Nrays(self):
712
+ '''
713
+ Nrays : int
714
+ Number of rays in angular discretisation used.
715
+ '''
716
+ try:
717
+ if self.muz is None:
718
+ raise AttributeError('Nrays not set, call atmos.rays or .quadrature first')
719
+ except AttributeError:
720
+ raise AttributeError('Nrays not set, call atmos.rays or .quadrature first')
721
+
722
+
723
+ return self.muz.shape[0]
724
+
725
+ def dimensioned_view(self):
726
+ '''
727
+ Returns a view over the contents of Layout reshaped so all data has
728
+ the correct (1/2/3D) dimensionality for the atmospheric model, as
729
+ these are all stored under a flat scheme.
730
+ '''
731
+ shape = self.structure.dimensioned_shape
732
+ atmos = copy(self)
733
+ atmos.structure = self.structure.dimensioned_view()
734
+ atmos.temperature = self.temperature.reshape(shape)
735
+ atmos.vturb = self.vturb.reshape(shape)
736
+ atmos.ne = self.ne.reshape(shape)
737
+ atmos.nHTot = self.nHTot.reshape(shape)
738
+ if self.B is not None:
739
+ atmos.B = self.B.reshape(shape)
740
+ atmos.chiB = self.chiB.reshape(shape)
741
+ atmos.gammaB = self.gammaB.reshape(shape)
742
+ return atmos
743
+
744
+ def unit_view(self):
745
+ '''
746
+ Returns a view over the contents of the Layout with the correct
747
+ `astropy.units`.
748
+ '''
749
+ atmos = copy(self)
750
+ atmos.structure = self.structure.unit_view()
751
+ atmos.temperature = self.temperature << u.K
752
+ atmos.vturb = self.vturb << u.m / u.s
753
+ atmos.ne = self.ne << u.m**(-3)
754
+ atmos.nHTot = self.nHTot << u.m**(-3)
755
+ if self.B is not None:
756
+ atmos.B = self.B << u.T
757
+ atmos.chiB = self.chiB << u.rad
758
+ atmos.gammaB = self.gammaB << u.rad
759
+ return atmos
760
+
761
+ def dimensioned_unit_view(self):
762
+ '''
763
+ Returns a view over the contents of Layout reshaped so all data has
764
+ the correct (1/2/3D) dimensionality for the atmospheric model, and
765
+ the correct `astropy.units`.
766
+ '''
767
+ atmos = self.dimensioned_view()
768
+ return atmos.unit_view()
769
+
770
+ @classmethod
771
+ def make_1d(
772
+ cls,
773
+ scale: ScaleType,
774
+ depthScale: np.ndarray,
775
+ temperature: np.ndarray,
776
+ vlos: Optional[np.ndarray]=None,
777
+ vturb: Optional[np.ndarray]=None,
778
+ ne: Optional[np.ndarray]=None,
779
+ hydrogenPops: Optional[np.ndarray]=None,
780
+ nHTot: Optional[np.ndarray]=None,
781
+ vx: Optional[np.ndarray]=None,
782
+ vy: Optional[np.ndarray]=None,
783
+ vz: Optional[np.ndarray]=None,
784
+ B: Optional[np.ndarray]=None,
785
+ gammaB: Optional[np.ndarray]=None,
786
+ chiB: Optional[np.ndarray]=None,
787
+ lowerBc: Optional[BoundaryCondition]=None,
788
+ upperBc: Optional[BoundaryCondition]=None,
789
+ convertScales: bool=True,
790
+ abundance: Optional[AtomicAbundance]=None,
791
+ logG: float=2.44,
792
+ Pgas: Optional[np.ndarray]=None,
793
+ Pe: Optional[np.ndarray]=None,
794
+ Ptop: Optional[float]=None,
795
+ PeTop: Optional[float]=None,
796
+ verbose: bool=False,
797
+ ):
798
+ '''
799
+ Constructor for 1D Atmosphere objects. Optionally will use an
800
+ equation of state (EOS) to estimate missing parameters.
801
+
802
+ If sufficient information is provided (i.e. all required parameters
803
+ and ne and (hydrogenPops or nHTot)) then the EOS is not invoked to
804
+ estimate any thermodynamic properties. If both of nHTot and
805
+ hydrogenPops are omitted, then the electron pressure will be used
806
+ with the Wittmann equation of state to estimate the mass density, and
807
+ the hydrogen number density will be inferred from this and the
808
+ abundances. If, instead, ne is omitted, then the mass density will be
809
+ used with the Wittmann EOS to estimate the electron pressure.
810
+ If both of these are omitted then the EOS will be used to estimate
811
+ both. If:
812
+
813
+ - Pgas is provided, then this gas pressure will define the
814
+ atmospheric stratification and will be used with the EOS.
815
+
816
+ - Pe is provided, then this electron pressure will define the
817
+ atmospheric stratification and will be used with the EOS.
818
+
819
+ - Ptop is provided, then this gas pressure at the top of the
820
+ atmosphere will be used with the log gravitational acceleration
821
+ logG, and the EOS to estimate the missing parameters assuming
822
+ hydrostatic equilibrium.
823
+
824
+ - PeTop is provided, then this electron pressure at the top of
825
+ the atmosphere will be used with the log gravitational
826
+ acceleration logG, and the EOS to estimate the missing parameters
827
+ assuming hydrostatic equilibrium.
828
+
829
+ - If all of Pgas, Pe, Ptop, PeTop are omitted then Ptop will be
830
+ estimated from the gas pressure in the FALC model at the
831
+ temperature at the top boundary. The hydrostatic reconstruction
832
+ will then continue as usual.
833
+
834
+ convertScales will substantially slow down this function due to the
835
+ slow calculation of background opacities used to compute tauRef. If
836
+ an atmosphere is constructed with a Geometric stratification, and an
837
+ estimate of tauRef is not required before running the main RT module,
838
+ then this can be set to False.
839
+ All of these parameters can be provided as astropy Quantities, and
840
+ will be converted in the constructor.
841
+
842
+ Parameters
843
+ ----------
844
+ scale : ScaleType
845
+ The type of stratification used along the z-axis.
846
+ depthScale : np.ndarray
847
+ The z-coordinates used along the chosen stratification. The
848
+ stratification is expected to start at the top of the atmosphere
849
+ (closest to the observer), and descend along the observer's line
850
+ of sight.
851
+ temperature : np.ndarray
852
+ Temperature structure of the atmosphere [K].
853
+ vlos : np.ndarray, optional
854
+ Alias for vz
855
+ Velocity structure of the atmosphere along z [m/s]. Default: 0 m/s everywhere
856
+ vturb : np.ndarray
857
+ Microturbulent velocity structure of the atmosphere [m/s]. Default: 0 m/s everywhere.
858
+ ne : np.ndarray
859
+ Electron density structure of the atmosphere [m-3].
860
+ hydrogenPops : np.ndarray, optional
861
+ Detailed (per level) hydrogen number density structure of the
862
+ atmosphere [m-3], 2D array [Nlevel, Nspace].
863
+ nHTot : np.ndarray, optional
864
+ Total hydrogen number density structure of the atmosphere [m-3]
865
+ vx : np.ndarray, optional
866
+ x-component of atmospheric velocity [m/s]. If specifying vx/vy then a 3D
867
+ quadrature will be needed.
868
+ vy : np.ndarray, optional
869
+ y-component of atmospheric velocity [m/s]. If specifying vx/vy then a 3D
870
+ quadrature will be needed.
871
+ vz : np.ndarray, optional
872
+ alias for vlos. [m/s]
873
+ B : np.ndarray, optional.
874
+ Magnetic field strength [T].
875
+ gammaB : np.ndarray, optional
876
+ Co-altitude of magnetic field vector [radians].
877
+ chiB : np.ndarray, optional
878
+ Azimuth of magnetic field vector (in x-y plane, from x) [radians].
879
+ lowerBc : BoundaryCondition, optional
880
+ Boundary condition for incoming radiation at the minimal z
881
+ coordinate (default: ThermalisedRadiation).
882
+ upperBc : BoundaryCondition, optional
883
+ Boundary condition for incoming radiation at the maximal z
884
+ coordinate (default: ZeroRadiation).
885
+ convertScales : bool, optional
886
+ Whether to automatically compute tauRef and cmass for an
887
+ atmosphere given in a stratification of m (default: True).
888
+ abundance: AtomicAbundance, optional
889
+ An instance of AtomicAbundance giving the abundances of each
890
+ atomic species in the given atmosphere, only used if the EOS is
891
+ invoked. (default: DefaultAtomicAbundance)
892
+ logG: float, optional
893
+ The log10 of the magnitude of gravitational acceleration [m/s2]
894
+ (default: 2.44).
895
+ Pgas: np.ndarray, optional
896
+ The gas pressure stratification of the atmosphere [Pa],
897
+ optionally used by the EOS.
898
+ Pe: np.ndarray, optional
899
+ The electron pressure stratification of the atmosphere [Pa],
900
+ optionally used by the EOS.
901
+ Ptop: np.ndarray, optional
902
+ The gas pressure at the top of the atmosphere [Pa], optionally
903
+ used by the EOS for a hydrostatic reconstruction.
904
+ Petop: np.ndarray, optional
905
+ The electron pressure at the top of the atmosphere [Pa],
906
+ optionally used by the EOS for a hydrostatic reconstruction.
907
+ verbose: bool, optional
908
+ Explain decisions made with the EOS to estimate missing
909
+ parameters (if invoked) through print calls (default: False).
910
+
911
+ Raises
912
+ ------
913
+ ValueError
914
+ if incorrect arguments or unable to construct estimate missing
915
+ parameters.
916
+ '''
917
+ if scale == ScaleType.Geometric:
918
+ depthScale = (depthScale << u.m).value
919
+ if np.any((depthScale[:-1] - depthScale[1:]) < 0.0):
920
+ raise ValueError("Geometric depth scale should be provided in decreasing height.")
921
+ elif scale == ScaleType.ColumnMass:
922
+ depthScale = (depthScale << u.kg / u.m**2).value
923
+ if np.any((depthScale[1:] - depthScale[:-1]) < 0.0):
924
+ raise ValueError("Column mass depth scale should be provided in increasing column mass.")
925
+
926
+ check_shape = lambda x, xName: check_shape_exception(x,
927
+ depthScale.shape[0], 1, xName)
928
+ temperature = (temperature << u.K).value
929
+ check_shape(temperature, 'temperature')
930
+ if vlos is None:
931
+ if vz is None:
932
+ vlos = np.zeros_like(temperature)
933
+ else:
934
+ vlos = vz
935
+ elif vz is not None:
936
+ raise ValueError("Cannot set both vlos and vz (they are aliases).")
937
+ vz = vlos
938
+ vlos = (vlos << u.m / u.s).value
939
+ check_shape(vlos, 'vlos')
940
+ if vturb is None:
941
+ vturb = np.zeros_like(temperature)
942
+ vturb = (vturb << u.m / u.s).value
943
+ check_shape(vturb, 'vturb')
944
+ if ne is not None:
945
+ ne = (ne << u.m**(-3)).value
946
+ check_shape(ne, 'ne')
947
+ if hydrogenPops is not None:
948
+ hydrogenPops = (hydrogenPops << u.m**(-3)).value
949
+ hydrogenPops = cast(np.ndarray, hydrogenPops)
950
+ if hydrogenPops.shape[1] != depthScale.shape[0]:
951
+ raise ValueError(f'Array hydrogenPops does not have the expected'
952
+ f' second dimension: {depthScale.shape[0]}'
953
+ f' (got: {hydrogenPops.shape[1]}).')
954
+ if nHTot is not None:
955
+ nHTot = (nHTot << u.m**(-3)).value
956
+ check_shape(nHTot, 'nHTot')
957
+ if vx is not None:
958
+ if vy is None:
959
+ raise ValueError("vx is set, vy must be also.")
960
+ check_shape(vx, "vx")
961
+ if vy is not None:
962
+ if vx is None:
963
+ raise ValueError("vy is set, vx must be also.")
964
+ check_shape(vy, "vy")
965
+
966
+ if B is not None:
967
+ B = (B << u.T).value
968
+ check_shape(B, 'B')
969
+ if gammaB is None or chiB is None:
970
+ raise ValueError('B is set, both gammaB and chiB must be also.')
971
+ if gammaB is not None:
972
+ gammaB = (gammaB << u.rad).value
973
+ check_shape(gammaB, 'gammaB')
974
+ if B is None or chiB is None:
975
+ raise ValueError('gammaB is set, both B and chiB must be also.')
976
+ if chiB is not None:
977
+ chiB = (chiB << u.rad).value
978
+ check_shape(chiB, 'chiB')
979
+ if gammaB is None or B is None:
980
+ raise ValueError('chiB is set, both B and gammaB must be also.')
981
+
982
+ if lowerBc is None:
983
+ lowerBc = ThermalisedRadiation()
984
+ elif isinstance(lowerBc, PeriodicRadiation):
985
+ raise ValueError('Cannot set periodic boundary conditions for 1D atmosphere')
986
+ if upperBc is None:
987
+ upperBc = ZeroRadiation()
988
+ elif isinstance(upperBc, PeriodicRadiation):
989
+ raise ValueError('Cannot set periodic boundary conditions for 1D atmosphere')
990
+
991
+ if scale != ScaleType.Geometric and not convertScales:
992
+ raise ValueError('Height scale must be provided if scale conversion is not applied')
993
+
994
+ if nHTot is None and hydrogenPops is not None:
995
+ nHTot = np.sum(hydrogenPops, axis=0)
996
+
997
+ if np.any(temperature < 2000):
998
+ # NOTE(cmo): Minimum value was decreased in NICOLE so should be safe
999
+ raise ValueError('Minimum temperature too low for EOS (< 2000 K)')
1000
+
1001
+ if abundance is None:
1002
+ abundance = DefaultAtomicAbundance
1003
+
1004
+ wittAbundances = np.array([abundance[e] for e in PeriodicTable.elements])
1005
+ eos = Wittmann(abund_init=wittAbundances)
1006
+
1007
+ Nspace = depthScale.shape[0]
1008
+ if nHTot is None and ne is not None:
1009
+ if verbose:
1010
+ print('Setting nHTot from electron pressure.')
1011
+ pe = (ne << u.Unit('m-3')).to('cm-3').value * cgs.BK * temperature
1012
+ rho = np.zeros(Nspace)
1013
+ for k in range(Nspace):
1014
+ rho[k] = eos.rho_from_pe(temperature[k], pe[k])
1015
+ nHTot = np.copy((rho << u.Unit('g cm-3')).to('kg m-3').value
1016
+ / (Const.Amu * abundance.massPerH))
1017
+ elif ne is None and nHTot is not None:
1018
+ if verbose:
1019
+ print('Setting ne from mass density.')
1020
+ rho = ((Const.Amu * abundance.massPerH * nHTot) << u.Unit('kg m-3')).to('g cm-3').value
1021
+ pe = np.zeros(Nspace)
1022
+ for k in range(Nspace):
1023
+ pe[k] = eos.pe_from_rho(temperature[k], rho[k])
1024
+ ne = np.copy(((pe / (cgs.BK * temperature)) << u.Unit('cm-3')).to('m-3').value)
1025
+ elif ne is None and nHTot is None:
1026
+ if Pgas is not None and Pgas.shape[0] != Nspace:
1027
+ raise ValueError('Dimensions of Pgas do not match atmospheric depth')
1028
+ if Pe is not None and Pe.shape[0] != Nspace:
1029
+ raise ValueError('Dimensions of Pe do not match atmospheric depth')
1030
+
1031
+ if Pgas is not None and Pe is None:
1032
+ if verbose:
1033
+ print('Setting ne, nHTot from provided gas pressure.')
1034
+ # Convert to cgs for eos
1035
+ pgas = (Pgas << u.Unit('Pa')).to('dyn cm-2').value
1036
+ pe = np.zeros(Nspace)
1037
+ rho = np.zeros(Nspace)
1038
+ for k in range(Nspace):
1039
+ pe[k] = eos.pe_from_pg(temperature[k], pgas[k])
1040
+ rho[k] = eos.rho_from_pg(temperature[k], pgas[k])
1041
+ elif Pe is not None and Pgas is None:
1042
+ if verbose:
1043
+ print('Setting ne, nHTot from provided electron pressure.')
1044
+ # Convert to cgs for eos
1045
+ pe = (Pe << u.Unit('Pa')).to('dyn cm-2').value
1046
+ pgas = np.zeros(Nspace)
1047
+ rho = np.zeros(Nspace)
1048
+ for k in range(Nspace):
1049
+ pgas[k] = eos.pg_from_pe(temperature[k], pe[k])
1050
+ rho[k] = eos.rho_from_pe(temperature[k], pe[k])
1051
+ elif Pgas is None and Pe is None:
1052
+ # Doing Hydrostatic Eq. based here on NICOLE implementation
1053
+ gravAcc = ((10**logG) << u.Unit('m s-2')).to('cm s-2').value
1054
+ Avog = 6.022045e23 # Avogadro's Number
1055
+ if Ptop is None and PeTop is not None:
1056
+ if verbose:
1057
+ print(('Setting ne, nHTot to hydrostatic equilibrium (logG=%f)'
1058
+ ' from provided top electron pressure.') % logG)
1059
+ PeTop = (PeTop << u.Unit("Pa")).to('dyn cm-2').value
1060
+ Ptop = eos.pg_from_pe(temperature[0], PeTop)
1061
+ elif Ptop is not None and PeTop is None:
1062
+ if verbose:
1063
+ print(('Setting ne, nHTot to hydrostatic equilibrium (logG=%f)'
1064
+ ' from provided top gas pressure.') % logG)
1065
+ Ptop = (Ptop << u.Unit("Pa")).to('dyn cm-2').value
1066
+ PeTop = eos.pe_from_pg(temperature[0], Ptop)
1067
+ elif Ptop is None and PeTop is None:
1068
+ if verbose:
1069
+ print(('Setting ne, nHTot to hydrostatic equilibrium (logG=%f)'
1070
+ ' from FALC gas pressure at upper boundary temperature.') % logG)
1071
+ Ptop = get_top_pressure(eos, temperature[0])
1072
+ PeTop = eos.pe_from_pg(temperature[0], Ptop)
1073
+ else:
1074
+ raise ValueError("Cannot set both Ptop and PeTop")
1075
+
1076
+ if scale == ScaleType.Tau500:
1077
+ tau = depthScale
1078
+ elif scale == ScaleType.Geometric:
1079
+ height = (depthScale << u.Unit('m')).to('cm').value
1080
+ else:
1081
+ cmass = (depthScale << u.Unit('kg m-2')).to('g cm-2').value
1082
+
1083
+ # NOTE(cmo): Compute HSE following the NICOLE method.
1084
+ rho = np.zeros(Nspace)
1085
+ chi_c = np.zeros(Nspace)
1086
+ pgas = np.zeros(Nspace)
1087
+ pe = np.zeros(Nspace)
1088
+ pgas[0] = Ptop
1089
+ pe[0] = PeTop
1090
+ chi_c[0] = eos.cont_opacity(temperature[0], pgas[0], pe[0],
1091
+ np.array([5000.0])).item()
1092
+ avg_mol_weight = lambda k: abundance.massPerH / (abundance.totalAbundance
1093
+ + pe[k] / pgas[k])
1094
+ rho[0] = Ptop * avg_mol_weight(0) / Avog / cgs.BK / temperature[0]
1095
+ chi_c[0] /= rho[0]
1096
+
1097
+ for k in range(1, Nspace):
1098
+ chi_c[k] = chi_c[k-1]
1099
+ rho[k] = rho[k-1]
1100
+ for it in range(200):
1101
+ if scale == ScaleType.Tau500:
1102
+ dtau = tau[k] - tau[k-1]
1103
+ pgas[k] = (pgas[k-1] + gravAcc * dtau
1104
+ / (0.5 * (chi_c[k-1] + chi_c[k])))
1105
+ elif scale == ScaleType.Geometric:
1106
+ pgas[k] = pgas[k-1] * np.exp(-gravAcc / Avog /
1107
+ cgs.BK * avg_mol_weight(k-1)
1108
+ * 0.5 * (1.0 / temperature[k-1]
1109
+ + 1.0 / temperature[k]) *
1110
+ (height[k] - height[k-1]))
1111
+ else:
1112
+ pgas[k] = gravAcc * cmass[k]
1113
+
1114
+ pe[k] = eos.pe_from_pg(temperature[k], pgas[k])
1115
+ prevChi = chi_c[k]
1116
+ chi_c[k] = eos.cont_opacity(temperature[k], pgas[k], pe[k],
1117
+ np.array([5000.0])).item()
1118
+ rho[k] = (pgas[k] * avg_mol_weight(k) / Avog /
1119
+ cgs.BK / temperature[k])
1120
+ chi_c[k] /= rho[k]
1121
+
1122
+ change = np.abs(prevChi - chi_c[k]) / (prevChi + chi_c[k])
1123
+ if change < 1e-5:
1124
+ break
1125
+ else:
1126
+ raise ConvergenceError(('No convergence in HSE at depth point %d, '
1127
+ 'last change %2.4e') % (k, change))
1128
+ nHTot = np.copy((rho << u.Unit('g cm-3')).to('kg m-3').value
1129
+ / (Const.Amu * abundance.massPerH))
1130
+ ne = np.copy(((pe / (cgs.BK * temperature)) << u.Unit('cm-3')).to('m-3').value)
1131
+
1132
+ # NOTE(cmo): Compute final pgas, pe from EOS that will be used for
1133
+ # background opacity.
1134
+ rhoSI = Const.Amu * abundance.massPerH * nHTot
1135
+ rho = (rhoSI << u.Unit('kg m-3')).to('g cm-3').value
1136
+ pgas = np.zeros_like(depthScale)
1137
+ pe = np.zeros_like(depthScale)
1138
+ for k in range(Nspace):
1139
+ pgas[k] = eos.pg_from_rho(temperature[k], rho[k])
1140
+ pe[k] = eos.pe_from_rho(temperature[k], rho[k])
1141
+
1142
+ chi_c = np.zeros_like(depthScale)
1143
+ for k in range(depthScale.shape[0]):
1144
+ chi_c[k] = eos.cont_opacity(temperature[k], pgas[k], pe[k],
1145
+ np.array([5000.0])).item()
1146
+ chi_c = (chi_c << u.Unit('cm')).to('m').value
1147
+
1148
+ # NOTE(cmo): We should now have a uniform minimum set of data (other
1149
+ # than the scale type), allowing us to simply convert between the
1150
+ # scales we do have!
1151
+ if convertScales:
1152
+ if scale == ScaleType.ColumnMass:
1153
+ height = np.zeros_like(depthScale)
1154
+ tau_ref = np.zeros_like(depthScale)
1155
+ cmass = depthScale
1156
+
1157
+ height[0] = 0.0
1158
+ tau_ref[0] = chi_c[0] / rhoSI[0] * cmass[0]
1159
+ for k in range(1, cmass.shape[0]):
1160
+ height[k] = height[k-1] - 2.0 * ((cmass[k] - cmass[k-1])
1161
+ / (rhoSI[k-1] + rhoSI[k]))
1162
+ tau_ref[k] = tau_ref[k-1] + 0.5 * ((chi_c[k-1] + chi_c[k])
1163
+ * (height[k-1] - height[k]))
1164
+
1165
+ hTau1 = np.interp(1.0, tau_ref, height)
1166
+ height -= hTau1
1167
+ elif scale == ScaleType.Geometric:
1168
+ cmass = np.zeros(Nspace)
1169
+ tau_ref = np.zeros(Nspace)
1170
+ height = depthScale
1171
+ nHTot = cast(np.ndarray, nHTot)
1172
+ ne = cast(np.ndarray, ne)
1173
+
1174
+ cmass[0] = ((nHTot[0] * abundance.massPerH + ne[0])
1175
+ * (Const.KBoltzmann * temperature[0] / 10**logG))
1176
+ tau_ref[0] = 0.5 * chi_c[0] * (height[0] - height[1])
1177
+ if tau_ref[0] > 1.0:
1178
+ tau_ref[0] = 0.0
1179
+
1180
+ for k in range(1, Nspace):
1181
+ cmass[k] = cmass[k-1] + 0.5 * ((rhoSI[k-1] + rhoSI[k])
1182
+ * (height[k-1] - height[k]))
1183
+ tau_ref[k] = tau_ref[k-1] + 0.5 * ((chi_c[k-1] + chi_c[k])
1184
+ * (height[k-1] - height[k]))
1185
+ elif scale == ScaleType.Tau500:
1186
+ cmass = np.zeros(Nspace)
1187
+ height = np.zeros(Nspace)
1188
+ tau_ref = depthScale
1189
+
1190
+ cmass[0] = (tau_ref[0] / chi_c[0]) * rhoSI[0]
1191
+ for k in range(1, Nspace):
1192
+ height[k] = height[k-1] - 2.0 * ((tau_ref[k] - tau_ref[k-1])
1193
+ / (chi_c[k-1] + chi_c[k]))
1194
+ cmass[k] = cmass[k-1] + 0.5 * ((chi_c[k-1] + chi_c[k])
1195
+ * (height[k-1] - height[k]))
1196
+
1197
+ hTau1 = np.interp(1.0, tau_ref, height)
1198
+ height -= hTau1
1199
+ else:
1200
+ raise ValueError("Other scales not handled yet")
1201
+
1202
+ stratifications: Optional[Stratifications] = Stratifications(
1203
+ cmass=cmass,
1204
+ tauRef=tau_ref)
1205
+
1206
+ else:
1207
+ stratifications = None
1208
+ height = depthScale
1209
+
1210
+ layout = Layout.make_1d(
1211
+ z=height,
1212
+ vx=vx,
1213
+ vy=vy,
1214
+ vz=vz,
1215
+ lowerBc=lowerBc,
1216
+ upperBc=upperBc,
1217
+ stratifications=stratifications
1218
+ )
1219
+ ne = cast(np.ndarray, ne)
1220
+ nHTot = cast(np.ndarray, nHTot)
1221
+ atmos = cls(structure=layout, temperature=temperature, vturb=vturb,
1222
+ ne=ne, nHTot=nHTot, B=B, gammaB=gammaB, chiB=chiB)
1223
+
1224
+ return atmos
1225
+
1226
+ @classmethod
1227
+ def make_2d(
1228
+ cls,
1229
+ height: np.ndarray,
1230
+ x: np.ndarray,
1231
+ temperature: np.ndarray,
1232
+ vx: Optional[np.ndarray]=None,
1233
+ vy: Optional[np.ndarray]=None,
1234
+ vz: Optional[np.ndarray]=None,
1235
+ vturb: Optional[np.ndarray]=None,
1236
+ ne: Optional[np.ndarray]=None,
1237
+ nHTot: Optional[np.ndarray]=None,
1238
+ B: Optional[np.ndarray]=None,
1239
+ gammaB: Optional[np.ndarray]=None,
1240
+ chiB: Optional[np.ndarray]=None,
1241
+ xUpperBc: Optional[BoundaryCondition]=None,
1242
+ xLowerBc: Optional[BoundaryCondition]=None,
1243
+ zUpperBc: Optional[BoundaryCondition]=None,
1244
+ zLowerBc: Optional[BoundaryCondition]=None,
1245
+ abundance: Optional[AtomicAbundance]=None,
1246
+ verbose=False
1247
+ ):
1248
+ '''
1249
+ Constructor for 2D Atmosphere objects.
1250
+
1251
+ No provision for estimating parameters using hydrostatic equilibrium
1252
+ is provided, but one of ne, or nHTot can be omitted and inferred by
1253
+ use of the Wittmann equation of state.
1254
+ The atmosphere must be defined on a geometric stratification.
1255
+ All atmospheric parameters are expected in a 2D [z, x] array.
1256
+
1257
+ Parameters
1258
+ ----------
1259
+ height : np.ndarray
1260
+ The z-coordinates of the atmospheric grid. The stratification is
1261
+ expected to start at the top of the atmosphere (closest to the
1262
+ observer), and descend along the observer's line of sight.
1263
+ x : np.ndarray
1264
+ The (horizontal) x-coordinates of the atmospheric grid.
1265
+ temperature : np.ndarray
1266
+ Temperature structure of the atmosphere [K].
1267
+ vx : np.ndarray, optional.
1268
+ x-component of the atmospheric velocity [m/s]. Default: 0 m/s.
1269
+ vy : np.ndarray, optional.
1270
+ y-component of the atmospheric velocity [m/s]. Not used by default,
1271
+ use a 3D quadrature if you need it.
1272
+ vz : np.ndarray, optional
1273
+ z-component of the atmospheric velocity [m/s]. Default: 0 m/s.
1274
+ vturb : np.ndarray
1275
+ Microturbulent velocity structure [m/s].
1276
+ ne : np.ndarray
1277
+ Electron density structure of the atmosphere [m-3].
1278
+ nHTot : np.ndarray, optional
1279
+ Total hydrogen number density structure of the atmosphere [m-3].
1280
+ B : np.ndarray, optional.
1281
+ Magnetic field strength [T].
1282
+ gammaB : np.ndarray, optional
1283
+ Inclination (co-altitude) of magnetic field vector to the z-axis
1284
+ [radians].
1285
+ chiB : np.ndarray, optional
1286
+ Azimuth of magnetic field vector (in x-y plane, from x) [radians].
1287
+ xLowerBc : BoundaryCondition, optional
1288
+ Boundary condition for incoming radiation at the minimal x
1289
+ coordinate (default: PeriodicRadiation).
1290
+ xUpperBc : BoundaryCondition, optional
1291
+ Boundary condition for incoming radiation at the maximal x
1292
+ coordinate (default: PeriodicRadiation).
1293
+ zLowerBc : BoundaryCondition, optional
1294
+ Boundary condition for incoming radiation at the minimal z
1295
+ coordinate (default: ThermalisedRadiation).
1296
+ zUpperBc : BoundaryCondition, optional
1297
+ Boundary condition for incoming radiation at the maximal z
1298
+ coordinate (default: ZeroRadiation).
1299
+ convertScales : bool, optional
1300
+ Whether to automatically compute tauRef and cmass for an
1301
+ atmosphere given in a stratification of m (default: True).
1302
+ abundance: AtomicAbundance, optional
1303
+ An instance of AtomicAbundance giving the abundances of each
1304
+ atomic species in the given atmosphere, only used if the EOS is
1305
+ invoked. (default: DefaultAtomicAbundance)
1306
+ verbose: bool, optional
1307
+ Explain decisions made with the EOS to estimate missing
1308
+ parameters (if invoked) through print calls (default: False).
1309
+
1310
+ Raises
1311
+ ------
1312
+ ValueError
1313
+ if incorrect arguments or unable to construct estimate missing
1314
+ parameters.
1315
+ '''
1316
+
1317
+ x = (x << u.m).value
1318
+ if np.any((x[1:] - x[:-1]) < 0.0):
1319
+ raise ValueError("x should be increasing with index (left -> right).")
1320
+ height = (height << u.m).value
1321
+ if np.any((height[:-1] - height[1:]) < 0.0):
1322
+ raise ValueError("Height should be decreasing with index (top -> bottom).")
1323
+ temperature = (temperature << u.K).value
1324
+ vx = (vx << u.m / u.s).value
1325
+ if vy is not None:
1326
+ vy = (vy << u.m / u.s).value
1327
+ vz = (vz << u.m / u.s).value
1328
+ vturb = (vturb << u.m / u.s).value
1329
+ if ne is not None:
1330
+ ne = (ne << u.m**(-3)).value
1331
+ if nHTot is not None:
1332
+ nHTot = (nHTot << u.m**(-3)).value
1333
+ if B is not None:
1334
+ B = (B << u.T).value
1335
+ B = cast(np.ndarray, B)
1336
+ flatB = view_flatten(B)
1337
+ else:
1338
+ flatB = None
1339
+
1340
+ if gammaB is not None:
1341
+ gammaB = (gammaB << u.rad).value
1342
+ gammaB = cast(np.ndarray, gammaB)
1343
+ flatGammaB = view_flatten(gammaB)
1344
+ else:
1345
+ flatGammaB = None
1346
+
1347
+ if chiB is not None:
1348
+ chiB = (chiB << u.rad).value
1349
+ chiB = cast(np.ndarray, chiB)
1350
+ flatChiB = view_flatten(chiB)
1351
+ else:
1352
+ flatChiB = None
1353
+
1354
+ if zLowerBc is None:
1355
+ zLowerBc = ThermalisedRadiation()
1356
+ elif isinstance(zLowerBc, PeriodicRadiation):
1357
+ raise ValueError('Cannot set periodic boundary conditions for z-axis.')
1358
+ if zUpperBc is None:
1359
+ zUpperBc = ZeroRadiation()
1360
+ elif isinstance(zUpperBc, PeriodicRadiation):
1361
+ raise ValueError('Cannot set periodic boundary conditions for z-axis.')
1362
+ if xUpperBc is None:
1363
+ xUpperBc = PeriodicRadiation()
1364
+ if xLowerBc is None:
1365
+ xLowerBc = PeriodicRadiation()
1366
+ if abundance is None:
1367
+
1368
+ abundance = DefaultAtomicAbundance
1369
+
1370
+ wittAbundances = np.array([abundance[e] for e in PeriodicTable.elements])
1371
+ eos = Wittmann(abund_init=wittAbundances)
1372
+
1373
+ flatHeight = view_flatten(height)
1374
+ flatTemperature = view_flatten(temperature)
1375
+ Nspace = flatHeight.shape[0]
1376
+ if nHTot is None and ne is not None:
1377
+ if verbose:
1378
+ print('Setting nHTot from electron pressure.')
1379
+ flatNe = view_flatten(ne)
1380
+ pe = (flatNe << u.Unit('m-3')).to('cm-3').value * cgs.BK * flatTemperature
1381
+ rho = np.zeros(Nspace)
1382
+ for k in range(Nspace):
1383
+ rho[k] = eos.rho_from_pe(flatTemperature[k], pe[k])
1384
+ nHTot = np.ascontiguousarray(
1385
+ (rho << u.Unit('g cm-3')).to('kg m-3').value
1386
+ / (Const.Amu * abundance.massPerH)
1387
+ )
1388
+ elif ne is None and nHTot is not None:
1389
+ if verbose:
1390
+ print('Setting ne from mass density.')
1391
+ flatNHTot = view_flatten(nHTot)
1392
+ rho = ((Const.Amu * abundance.massPerH * flatNHTot) << u.Unit('kg m-3')).to('g cm-3').value
1393
+ pe = np.zeros(Nspace)
1394
+ for k in range(Nspace):
1395
+ pe[k] = eos.pe_from_rho(flatTemperature[k], rho[k])
1396
+ ne = np.ascontiguousarray(((pe / (cgs.BK * flatTemperature)) << u.Unit('cm-3')).to('m-3').value)
1397
+ elif ne is None and nHTot is None:
1398
+ raise ValueError('Cannot omit both ne and nHTot (currently).')
1399
+ flatX = view_flatten(x)
1400
+ nHTot = cast(np.ndarray, nHTot)
1401
+ flatNHTot = view_flatten(nHTot)
1402
+ ne = cast(np.ndarray, ne)
1403
+ flatNe = view_flatten(ne)
1404
+ flatVx = view_flatten(vx)
1405
+ flatVy = None if vy is None else view_flatten(vy)
1406
+ flatVz = view_flatten(vz)
1407
+ flatVturb = view_flatten(vturb)
1408
+
1409
+ layout = Layout.make_2d(
1410
+ x=flatX,
1411
+ z=flatHeight,
1412
+ vx=flatVx,
1413
+ vy=flatVy,
1414
+ vz=flatVz,
1415
+ xLowerBc=xLowerBc,
1416
+ xUpperBc=xUpperBc,
1417
+ zLowerBc=zLowerBc,
1418
+ zUpperBc=zUpperBc,
1419
+ stratifications=None,
1420
+ )
1421
+
1422
+ atmos = cls(structure=layout, temperature=flatTemperature,
1423
+ vturb=flatVturb, ne=flatNe, nHTot=flatNHTot, B=flatB,
1424
+ gammaB=flatGammaB, chiB=flatChiB)
1425
+ return atmos
1426
+
1427
+
1428
+ def quadrature(
1429
+ self,
1430
+ Nrays: Optional[int]=None,
1431
+ mu: Optional[Sequence[float]]=None,
1432
+ wmu: Optional[Sequence[float]]=None,
1433
+ force3d: bool=False,
1434
+ ):
1435
+ '''
1436
+ Compute the angular quadrature for solving the RTE and Kinetic
1437
+ Equilibrium in a given atmosphere.
1438
+
1439
+ Procedure varies with dimensionality.
1440
+
1441
+ By convention muz is always positive, as the direction on this axis
1442
+ is determined by the toObs term that is used internally to the formal
1443
+ solver.
1444
+
1445
+ 1D:
1446
+ If a number of rays is given (typically 3 or 5), then the
1447
+ Gauss-Legendre quadrature for this set is used.
1448
+ If mu and wmu are instead given then these will be validated and
1449
+ used.
1450
+
1451
+ 2+D:
1452
+ If the number of rays selected is in the list of near optimal
1453
+ quadratures for unpolarised radiation provided by Stepan et al
1454
+ 2020 (A&A, 646 A24), then this is used. Otherwise an exception is
1455
+ raised.
1456
+
1457
+ The available quadratures are:
1458
+
1459
+ +--------+-------+
1460
+ | Points | Order |
1461
+ +========+=======+
1462
+ | 1 | 3 |
1463
+ +--------+-------+
1464
+ | 3 | 7 |
1465
+ +--------+-------+
1466
+ | 6 | 9 |
1467
+ +--------+-------+
1468
+ | 7 | 11 |
1469
+ +--------+-------+
1470
+ | 10 | 13 |
1471
+ +--------+-------+
1472
+ | 11 | 15 |
1473
+ +--------+-------+
1474
+
1475
+ Parameters
1476
+ ----------
1477
+ Nrays : int, optional
1478
+ The number of rays to use in the quadrature (per octant). See notes
1479
+ above.
1480
+ mu : sequence of float, optional
1481
+ The cosine of the angle made between the between each of the set
1482
+ of rays and the z axis, only used in 1D.
1483
+ wmu : sequence of float, optional
1484
+ The integration weights for each mu, must be provided if mu is provided.
1485
+ force3d : bool, optional
1486
+ Force the use of a 3D quadrature. Default: False.
1487
+
1488
+ Raises
1489
+ ------
1490
+ ValueError
1491
+ on incorrect input.
1492
+ '''
1493
+
1494
+ n_dim_effective = self.Ndim if not force3d else 3
1495
+ # NOTE(cmo): Catch the case where we need a 3d quadrature in a less than 3d atmosphere.
1496
+ if self.structure.vx.size > 0 and self.structure.vy.size > 0:
1497
+ n_dim_effective = 3
1498
+ if n_dim_effective == 1:
1499
+ if Nrays is not None and mu is None:
1500
+ if Nrays >= 1:
1501
+ x, w = leggauss(Nrays)
1502
+ mid, halfWidth = 0.5, 0.5
1503
+ x = mid + halfWidth * x
1504
+ w *= halfWidth
1505
+
1506
+ # NOTE(cmo): muz is in [0, 1], i.e. the upward directed
1507
+ # rays. The downward rays are handled by up/down/
1508
+ self.muz = x
1509
+ self.wmu = w
1510
+ else:
1511
+ raise ValueError('Unsupported Nrays=%d' % Nrays)
1512
+ elif Nrays is not None and mu is not None:
1513
+ if wmu is None:
1514
+ raise ValueError('Must provide wmu when providing mu')
1515
+ if Nrays != len(mu):
1516
+ raise ValueError('mu must be Nrays long if Nrays is provided')
1517
+ if len(mu) != len(wmu):
1518
+ raise ValueError('mu and wmu must be the same shape')
1519
+
1520
+ self.muz = np.array(mu, dtype=np.float64)
1521
+ self.wmu = np.array(wmu, dtype=np.float64)
1522
+
1523
+ self.muy = np.zeros_like(self.muz)
1524
+ self.mux = np.sqrt(1.0 - self.muz**2)
1525
+ else:
1526
+ with open(get_data_path() + 'Quadratures.pickle', 'rb') as pkl:
1527
+ quads = pickle.load(pkl)
1528
+
1529
+ rays = {int(q.split('n')[1]): q for q in quads}
1530
+ if Nrays not in rays:
1531
+ raise ValueError('For multidimensional cases Nrays must be in %s' % repr(rays))
1532
+
1533
+ quad = quads[rays[Nrays]]
1534
+
1535
+ if n_dim_effective == 2:
1536
+ Nrays *= 2
1537
+ theta = np.deg2rad(quad[:, 1])
1538
+ chi = np.deg2rad(quad[:, 2])
1539
+ # polar coords:
1540
+ # x = sin theta cos chi
1541
+ # y = sin theta sin chi
1542
+ # z = cos theta
1543
+ # Fill first half then flip in x-z plane. This solves for 2 octants.
1544
+ # As always, we keep muz always positive in the atmosphere definition.
1545
+ self.mux = np.zeros(Nrays)
1546
+ self.mux[:Nrays // 2] = np.sin(theta) * np.cos(chi)
1547
+ self.mux[Nrays // 2:] = -np.sin(theta) * np.cos(chi)
1548
+ self.muz = np.zeros(Nrays)
1549
+ self.muz[:Nrays // 2] = np.cos(theta)
1550
+ self.muz[Nrays // 2:] = np.cos(theta)
1551
+ self.wmu = np.zeros(Nrays)
1552
+ self.wmu[:Nrays // 2] = quad[:, 0]
1553
+ self.wmu[Nrays // 2:] = quad[:, 0]
1554
+ self.wmu /= np.sum(self.wmu)
1555
+ self.muy = np.sqrt(1.0 - (self.mux**2 + self.muz**2))
1556
+ else:
1557
+ # n_dim_effective == 3
1558
+ Nrays *= 4
1559
+ theta = np.deg2rad(quad[:, 1])
1560
+ chi = np.deg2rad(quad[:, 2])
1561
+ # polar coords:
1562
+ # x = sin theta cos chi
1563
+ # y = sin theta sin chi
1564
+ # z = cos theta
1565
+ # Fill first quarter then flip in x-z plane for second quarter.
1566
+ # Then flip in y-z plane for second half. This solves for 4 octants.
1567
+ # As always, we keep muz always positive in the atmosphere definition.
1568
+ self.mux = np.zeros(Nrays)
1569
+ self.mux[:Nrays // 4] = np.sin(theta) * np.cos(chi)
1570
+ self.mux[Nrays // 4:Nrays // 2] = -np.sin(theta) * np.cos(chi)
1571
+ self.mux[Nrays // 2:] = -self.mux[:Nrays // 2]
1572
+ self.muy = np.zeros(Nrays)
1573
+ self.muy[:Nrays // 4] = np.sin(theta) * np.sin(chi)
1574
+ self.muy[Nrays // 4:Nrays // 2] = np.sin(theta) * np.sin(chi)
1575
+ self.muy[Nrays // 2:] = -self.muy[:Nrays // 2]
1576
+
1577
+ self.wmu = np.zeros(Nrays)
1578
+ self.wmu[:Nrays // 4] = quad[:, 0]
1579
+ self.wmu[Nrays // 4:Nrays // 2] = quad[:, 0]
1580
+ self.wmu[Nrays // 2:] = self.wmu[:Nrays // 2]
1581
+
1582
+ self.wmu /= np.sum(self.wmu)
1583
+ self.muz = np.sqrt(1.0 - (self.mux**2 + self.muy**2))
1584
+
1585
+
1586
+ self.configure_bcs()
1587
+
1588
+
1589
+ def rays(self, muz: Union[float, Sequence[float]],
1590
+ mux: Optional[Union[float, Sequence[float]]]=None,
1591
+ muy: Optional[Union[float, Sequence[float]]]=None,
1592
+ wmu: Optional[Union[float, Sequence[float]]]=None,
1593
+ upOnly: bool=False):
1594
+ '''
1595
+ Set up the rays on the Atmosphere for computing the intensity in a
1596
+ particular direction (or set of directions).
1597
+
1598
+ If only the z angle is set then the ray is assumed in the x-z plane.
1599
+ If either muz or muy is omitted then this angle is inferred by
1600
+ normalisation of the projection.
1601
+
1602
+ By convention muz is always positive, as the direction on this axis
1603
+ is determined by the toObs term that is used internally to the formal
1604
+ solver.
1605
+
1606
+ Parameters
1607
+ ----------
1608
+ muz : float or sequence of float, optional
1609
+ The angular projections along the z axis.
1610
+ mux : float or sequence of float, optional
1611
+ The angular projections along the x axis.
1612
+ muy : float or sequence of float, optional
1613
+ The angular projections along the y axis.
1614
+ wmu : float or sequence of float, optional
1615
+ The integration weights for the given ray if J is to be
1616
+ integrated for angle set.
1617
+ upOnly : bool, optional
1618
+ Whether to only configure boundary conditions for up-only rays.
1619
+ (default: False)
1620
+
1621
+ Raises
1622
+ ------
1623
+ ValueError
1624
+ if the angular projections or integration weights are incorrectly
1625
+ normalised.
1626
+ '''
1627
+
1628
+ if isinstance(muz, numbers.Real):
1629
+ muz = [float(muz)]
1630
+ if isinstance(mux, numbers.Real):
1631
+ mux = [float(mux)]
1632
+ if isinstance(muy, numbers.Real):
1633
+ muy = [float(muy)]
1634
+ if isinstance(wmu, numbers.Real):
1635
+ wmu = [float(wmu)]
1636
+
1637
+ if mux is None and muy is None:
1638
+ self.muz = np.array(muz, dtype=np.float64)
1639
+ self.wmu = np.zeros_like(self.muz)
1640
+ self.muy = np.zeros_like(self.muz)
1641
+ self.mux = np.sqrt(1.0 - self.muz**2)
1642
+ elif muy is None:
1643
+ self.muz = np.array(muz, dtype=np.float64)
1644
+ self.wmu = np.zeros_like(self.muz)
1645
+ self.mux = np.array(mux, dtype=np.float64)
1646
+ self.muy = np.sqrt(1.0 - (self.muz**2 + self.mux**2))
1647
+ elif mux is None:
1648
+ self.muz = np.array(muz, dtype=np.float64)
1649
+ self.wmu = np.zeros_like(self.muz)
1650
+ self.muy = np.array(muy, dtype=np.float64)
1651
+ self.mux = np.sqrt(1.0 - (self.muz**2 + self.muy**2))
1652
+ else:
1653
+ self.muz = np.array(muz, dtype=np.float64)
1654
+ self.mux = np.array(mux, dtype=np.float64)
1655
+ self.muy = np.array(muy, dtype=np.float64)
1656
+ self.wmu = np.zeros_like(muz)
1657
+
1658
+ if not np.allclose(self.muz**2 + self.mux**2 + self.muy**2, 1):
1659
+ raise ValueError('mux**2 + muy**2 + muz**2 != 1.0')
1660
+
1661
+ if not np.all(self.muz > 0):
1662
+ raise ValueError('muz must be > 0')
1663
+
1664
+ if wmu is not None:
1665
+ self.wmu = np.array(wmu, dtype=np.float64)
1666
+
1667
+ if not np.isclose(self.wmu.sum(), 1.0):
1668
+ raise ValueError('sum of wmus is not 1.0')
1669
+
1670
+ self.configure_bcs(upOnly=upOnly)
1671
+
1672
+ def configure_bcs(self, upOnly: bool=False):
1673
+ '''
1674
+ Configure the required angular information for all boundary
1675
+ conditions on the model.
1676
+
1677
+ Parameters
1678
+ ----------
1679
+ upOnly : bool, optional
1680
+ Whether to only configure boundary conditions for up-going rays.
1681
+ (default: False)
1682
+ '''
1683
+
1684
+ # NOTE(cmo): We always have z-bcs
1685
+ # For zLowerBc, muz is positive, and we have all mux, muz
1686
+ mux, muy, muz = self.mux, self.muy, self.muz
1687
+ # NOTE(cmo): indexVector is of shape (mu, toObs) to allow the core to
1688
+ # easily destructure the blob that will be handed to it from
1689
+ # compute_bc.
1690
+ indexVector = np.ones((self.mux.shape[0], 2), dtype=np.int32) * -1
1691
+ indexVector[:, 1] = np.arange(mux.shape[0])
1692
+ self.zLowerBc.set_required_angles(mux, muy, muz, indexVector)
1693
+
1694
+ indexVector = np.ones((mux.shape[0], 2), dtype=np.int32) * -1
1695
+ if not upOnly:
1696
+ indexVector[:, 0] = np.arange(mux.shape[0])
1697
+ self.zUpperBc.set_required_angles(-mux, -muy, -muz, indexVector)
1698
+
1699
+ toObsRange = [0, 1]
1700
+ if upOnly:
1701
+ toObsRange = [1]
1702
+
1703
+ # NOTE(cmo): If 2+D we have x-bcs too
1704
+ # xLowerBc has all muz and all mux > 0
1705
+ mux, muy, muz = [], [], []
1706
+ indexVector = np.ones((self.mux.shape[0], 2), dtype=np.int32) * -1
1707
+ count = 0
1708
+ musDone = np.zeros(self.muz.shape[0], dtype=np.bool_)
1709
+ for mu in range(self.muz.shape[0]):
1710
+ for equalMu in np.argwhere(np.abs(self.muz) == self.muz[mu]).reshape(-1)[::-1]:
1711
+ if musDone[equalMu]:
1712
+ continue
1713
+ musDone[equalMu] = True
1714
+
1715
+ for toObsI in toObsRange:
1716
+ sign = [-1, 1][toObsI]
1717
+ sMux = sign * self.mux[equalMu]
1718
+ if sMux > 0:
1719
+ mux.append(sMux)
1720
+ muy.append(sign * self.muy[equalMu])
1721
+ muz.append(sign * self.muz[equalMu])
1722
+ indexVector[equalMu, toObsI] = count
1723
+ count += 1
1724
+ if np.all(musDone):
1725
+ break
1726
+
1727
+ mux = np.array(mux)
1728
+ muy = np.array(muy)
1729
+ muz = np.array(muz)
1730
+ self.xLowerBc.set_required_angles(mux, muy, muz, indexVector)
1731
+
1732
+ mux, muy, muz = [], [], []
1733
+ indexVector = np.ones((self.mux.shape[0], 2), dtype=np.int32) * -1
1734
+ count = 0
1735
+ musDone = np.zeros(self.muz.shape[0], dtype=np.bool_)
1736
+ for mu in range(self.muz.shape[0]):
1737
+ for equalMu in np.argwhere(np.abs(self.muz) == self.muz[mu]).reshape(-1):
1738
+ if musDone[equalMu]:
1739
+ continue
1740
+ musDone[equalMu] = True
1741
+
1742
+ for toObsI in toObsRange:
1743
+ sign = [-1, 1][toObsI]
1744
+ sMux = sign * self.mux[equalMu]
1745
+ if sMux < 0:
1746
+ mux.append(sMux)
1747
+ muy.append(sign * self.muy[equalMu])
1748
+ muz.append(sign * self.muz[equalMu])
1749
+ indexVector[equalMu, toObsI] = count
1750
+ count += 1
1751
+ if np.all(musDone):
1752
+ break
1753
+
1754
+ mux = np.array(mux)
1755
+ muy = np.array(muy)
1756
+ muz = np.array(muz)
1757
+ self.xUpperBc.set_required_angles(mux, muy, muz, indexVector)
1758
+
1759
+ self.yLowerBc.set_required_angles(np.zeros((0)), np.zeros((0)), np.zeros((0)),
1760
+ np.ones((self.mux.shape[0], 2),
1761
+ dtype=np.int32) * -1)
1762
+ self.yUpperBc.set_required_angles(np.zeros((0)), np.zeros((0)), np.zeros((0)),
1763
+ np.ones((self.mux.shape[0], 2),
1764
+ dtype=np.int32) * -1)
1765
+
1766
+ if self.Ndim > 2:
1767
+ raise ValueError('Only <= 2D atmospheres supported currently.')