solrat 0.9.6__tar.gz

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 (64) hide show
  1. solrat-0.9.6/LICENSE.md +1 -0
  2. solrat-0.9.6/PKG-INFO +54 -0
  3. solrat-0.9.6/README.md +33 -0
  4. solrat-0.9.6/pyproject.toml +35 -0
  5. solrat-0.9.6/setup.cfg +4 -0
  6. solrat-0.9.6/solrat/__init__.py +0 -0
  7. solrat-0.9.6/solrat/about.py +1 -0
  8. solrat-0.9.6/solrat/common/__init__.py +0 -0
  9. solrat-0.9.6/solrat/common/constants.py +19 -0
  10. solrat-0.9.6/solrat/common/functions.py +44 -0
  11. solrat-0.9.6/solrat/common/rotations.py +120 -0
  12. solrat-0.9.6/solrat/common/voigt_profile.py +45 -0
  13. solrat-0.9.6/solrat/common/wigner_3j_6j_9j.py +216 -0
  14. solrat-0.9.6/solrat/engine/__init__.py +0 -0
  15. solrat-0.9.6/solrat/engine/functions/__init__.py +0 -0
  16. solrat-0.9.6/solrat/engine/functions/decorators.py +248 -0
  17. solrat-0.9.6/solrat/engine/functions/general.py +55 -0
  18. solrat-0.9.6/solrat/engine/functions/looping.py +47 -0
  19. solrat-0.9.6/solrat/engine/functions/special.py +13 -0
  20. solrat-0.9.6/solrat/engine/generators/__init__.py +0 -0
  21. solrat-0.9.6/solrat/engine/generators/merge_frame.py +302 -0
  22. solrat-0.9.6/solrat/engine/generators/merge_loopers.py +450 -0
  23. solrat-0.9.6/solrat/engine/generators/multiply.py +67 -0
  24. solrat-0.9.6/solrat/engine/generators/nested_loops.py +35 -0
  25. solrat-0.9.6/solrat/engine/generators/summate.py +48 -0
  26. solrat-0.9.6/solrat/gui/__init__.py +0 -0
  27. solrat-0.9.6/solrat/gui/plots/__init__.py +0 -0
  28. solrat-0.9.6/solrat/gui/plots/plot_stokes_profiles.py +229 -0
  29. solrat-0.9.6/solrat/multi_term_atom/__init__.py +0 -0
  30. solrat-0.9.6/solrat/multi_term_atom/atmosphere/__init__.py +0 -0
  31. solrat-0.9.6/solrat/multi_term_atom/atmosphere/constant_property_slab.py +151 -0
  32. solrat-0.9.6/solrat/multi_term_atom/atmosphere/multi_slab_atmosphere.py +27 -0
  33. solrat-0.9.6/solrat/multi_term_atom/atomic_data/FeI.py +66 -0
  34. solrat-0.9.6/solrat/multi_term_atom/atomic_data/HI.py +98 -0
  35. solrat-0.9.6/solrat/multi_term_atom/atomic_data/HeI.py +190 -0
  36. solrat-0.9.6/solrat/multi_term_atom/atomic_data/HeI_precomputed/__init__.py +0 -0
  37. solrat-0.9.6/solrat/multi_term_atom/atomic_data/MnI.py +56 -0
  38. solrat-0.9.6/solrat/multi_term_atom/atomic_data/NiI.py +42 -0
  39. solrat-0.9.6/solrat/multi_term_atom/atomic_data/__init__.py +0 -0
  40. solrat-0.9.6/solrat/multi_term_atom/atomic_data/mock.py +50 -0
  41. solrat-0.9.6/solrat/multi_term_atom/legacy/__init__.py +0 -0
  42. solrat-0.9.6/solrat/multi_term_atom/legacy/radiative_transfer_equations_legacy.py +1012 -0
  43. solrat-0.9.6/solrat/multi_term_atom/legacy/statistical_equilibrium_equations_legacy.py +596 -0
  44. solrat-0.9.6/solrat/multi_term_atom/object/__init__.py +0 -0
  45. solrat-0.9.6/solrat/multi_term_atom/object/angles.py +18 -0
  46. solrat-0.9.6/solrat/multi_term_atom/object/atmosphere_parameters.py +34 -0
  47. solrat-0.9.6/solrat/multi_term_atom/object/multi_term_atom_context.py +48 -0
  48. solrat-0.9.6/solrat/multi_term_atom/object/radiation_tensor.py +198 -0
  49. solrat-0.9.6/solrat/multi_term_atom/object/radiative_transfer_coefficients.py +140 -0
  50. solrat-0.9.6/solrat/multi_term_atom/object/rho_matrix_builder.py +137 -0
  51. solrat-0.9.6/solrat/multi_term_atom/object/stokes.py +55 -0
  52. solrat-0.9.6/solrat/multi_term_atom/physics/__init__.py +0 -0
  53. solrat-0.9.6/solrat/multi_term_atom/physics/einstein_coefficients.py +18 -0
  54. solrat-0.9.6/solrat/multi_term_atom/physics/paschen_back.py +161 -0
  55. solrat-0.9.6/solrat/multi_term_atom/radiative_transfer_equations.py +564 -0
  56. solrat-0.9.6/solrat/multi_term_atom/statistical_equilibrium_equations.py +917 -0
  57. solrat-0.9.6/solrat/multi_term_atom/terms_levels_transitions/__init__.py +0 -0
  58. solrat-0.9.6/solrat/multi_term_atom/terms_levels_transitions/level_registry.py +199 -0
  59. solrat-0.9.6/solrat/multi_term_atom/terms_levels_transitions/transition_registry.py +130 -0
  60. solrat-0.9.6/solrat.egg-info/PKG-INFO +54 -0
  61. solrat-0.9.6/solrat.egg-info/SOURCES.txt +62 -0
  62. solrat-0.9.6/solrat.egg-info/dependency_links.txt +1 -0
  63. solrat-0.9.6/solrat.egg-info/requires.txt +6 -0
  64. solrat-0.9.6/solrat.egg-info/top_level.txt +1 -0
@@ -0,0 +1 @@
1
+ Copyright (C) 2023, Ivan I. Yakovkin. Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL Ivan I. Yakovkin BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. Except as contained in this notice, the name of Ivan I. Yakovkin shall not be used in advertising or otherwise to promote the sale, use or other dealings in this Software without prior written authorization from Ivan I. Yakovkin.
solrat-0.9.6/PKG-INFO ADDED
@@ -0,0 +1,54 @@
1
+ Metadata-Version: 2.4
2
+ Name: solrat
3
+ Version: 0.9.6
4
+ Summary: Solar Radiative Transfer non-LTE forward-modeling code.
5
+ Author: Ivan I. Yakovkin
6
+ License: MIT
7
+ Project-URL: Homepage, https://github.com/yakovkinii/solrat
8
+ Classifier: Programming Language :: Python :: 3
9
+ Classifier: License :: OSI Approved :: MIT License
10
+ Classifier: Operating System :: OS Independent
11
+ Requires-Python: >=3.8
12
+ Description-Content-Type: text/markdown
13
+ License-File: LICENSE.md
14
+ Requires-Dist: numpy
15
+ Requires-Dist: pandas
16
+ Requires-Dist: matplotlib
17
+ Requires-Dist: tqdm
18
+ Requires-Dist: yatools
19
+ Requires-Dist: sympy
20
+ Dynamic: license-file
21
+
22
+ <H1>SolRaT</H1>
23
+
24
+ SolRaT (Solar Radiative Transfer) is a forward modeling code for non-LTE (and optionally LTE) transfer of
25
+ radiation in stellar atmospheres.
26
+
27
+ The code implements the multi-term atom model, described in
28
+ `Landi Degl’Innocenti, E., & Landolfi, M. 2004, Polarization in Spectral Lines (Dordrecht: Kluwer)`.
29
+ SolRaT supports atomic level polarization, arbitrary magnetic fields (intermediate Paschen-Back effect),
30
+ Hanle effect and many other features.
31
+
32
+ The code is written in python, currently tested on Windows and Ubuntu Linux.
33
+ The code is expected to work on all systems that fully support python 3.11.
34
+ SolRaT is currently in beta testing. Journal article and detailed documentation are pending.
35
+ Until then, if SolRaT has found use in your research, please cite it as
36
+ ```
37
+ Yakovkin I. I. SolRaT (2023) [computer software]. Retrieved from https://www.yakovkinii.com/solrat/
38
+ ```
39
+
40
+ How to run:
41
+ ```bash
42
+ git clone https://github.com/yakovkinii/SolRaT.git
43
+ pip install -r requirements.txt
44
+ python ./run_all_tests.py
45
+ ```
46
+
47
+ Some examples of how to use SolRaT are available in the `_demos` directory.
48
+
49
+ Keywords:
50
+ Non-LTE, Stokes Profiles, Inversion, Synthesis, Paschen-Back, Hanle, Zeeman,
51
+ Magnetic Fields, Sun, Solar Atmosphere, Radiative Transfer, Polarization,
52
+ Spectral Lines, Two-Term Atom Model
53
+
54
+ Copyright (2023) Ivan I. Yakovkin
solrat-0.9.6/README.md ADDED
@@ -0,0 +1,33 @@
1
+ <H1>SolRaT</H1>
2
+
3
+ SolRaT (Solar Radiative Transfer) is a forward modeling code for non-LTE (and optionally LTE) transfer of
4
+ radiation in stellar atmospheres.
5
+
6
+ The code implements the multi-term atom model, described in
7
+ `Landi Degl’Innocenti, E., & Landolfi, M. 2004, Polarization in Spectral Lines (Dordrecht: Kluwer)`.
8
+ SolRaT supports atomic level polarization, arbitrary magnetic fields (intermediate Paschen-Back effect),
9
+ Hanle effect and many other features.
10
+
11
+ The code is written in python, currently tested on Windows and Ubuntu Linux.
12
+ The code is expected to work on all systems that fully support python 3.11.
13
+ SolRaT is currently in beta testing. Journal article and detailed documentation are pending.
14
+ Until then, if SolRaT has found use in your research, please cite it as
15
+ ```
16
+ Yakovkin I. I. SolRaT (2023) [computer software]. Retrieved from https://www.yakovkinii.com/solrat/
17
+ ```
18
+
19
+ How to run:
20
+ ```bash
21
+ git clone https://github.com/yakovkinii/SolRaT.git
22
+ pip install -r requirements.txt
23
+ python ./run_all_tests.py
24
+ ```
25
+
26
+ Some examples of how to use SolRaT are available in the `_demos` directory.
27
+
28
+ Keywords:
29
+ Non-LTE, Stokes Profiles, Inversion, Synthesis, Paschen-Back, Hanle, Zeeman,
30
+ Magnetic Fields, Sun, Solar Atmosphere, Radiative Transfer, Polarization,
31
+ Spectral Lines, Two-Term Atom Model
32
+
33
+ Copyright (2023) Ivan I. Yakovkin
@@ -0,0 +1,35 @@
1
+ [build-system]
2
+ requires = ["setuptools>=61.0", "wheel"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "solrat"
7
+ dynamic = ["version"]
8
+ description = "Solar Radiative Transfer non-LTE forward-modeling code."
9
+ authors = [{name = "Ivan I. Yakovkin"}]
10
+ license = {text = "MIT"}
11
+ readme = "README.md"
12
+ requires-python = ">=3.8"
13
+ dependencies = [
14
+ "numpy",
15
+ "pandas",
16
+ "matplotlib",
17
+ "tqdm",
18
+ "yatools",
19
+ "sympy",
20
+ ]
21
+ classifiers = [
22
+ "Programming Language :: Python :: 3",
23
+ "License :: OSI Approved :: MIT License",
24
+ "Operating System :: OS Independent",
25
+ ]
26
+
27
+ [project.urls]
28
+ Homepage = "https://github.com/yakovkinii/solrat"
29
+
30
+ [tool.setuptools.dynamic]
31
+ version = {attr = "solrat.about.version"}
32
+
33
+ [tool.setuptools.packages.find]
34
+ where = ["."]
35
+ include = ["solrat*"]
solrat-0.9.6/setup.cfg ADDED
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
File without changes
@@ -0,0 +1 @@
1
+ version = "0.9.6"
File without changes
@@ -0,0 +1,19 @@
1
+ """
2
+ TODO
3
+ TODO This file needs improved documentation.
4
+ TODO
5
+ """
6
+
7
+ import numpy as np
8
+
9
+ sqrt2 = np.sqrt(2)
10
+ sqrt3 = np.sqrt(3)
11
+ sqrt_pi = np.sqrt(np.pi)
12
+
13
+ h_erg_s = 6.626196e-27 # erg s
14
+ c_cm_sm1 = 2.99792458e10 # cm/s
15
+ kB_erg_Km1 = 1.380658e-16 # erg/K
16
+ # e0_cgs = 0.480321e-9 # StatCoulomb
17
+ # me_g = 9.109390e-28 # g
18
+ atomic_mass_unit_g = 1.660539e-24 # g
19
+ mu0_erg_gaussm1 = 9.274e-21 # Bohr magneton in erg/G
@@ -0,0 +1,44 @@
1
+ """
2
+ TODO
3
+ TODO This file needs improved documentation.
4
+ TODO
5
+ """
6
+
7
+ from typing import Union
8
+
9
+ import numpy as np
10
+ from numpy import exp
11
+
12
+ from solrat.common.constants import c_cm_sm1, h_erg_s, kB_erg_Km1, mu0_erg_gaussm1
13
+
14
+
15
+ def get_planck_BP(nu_sm1: Union[float, np.ndarray], T_K: float) -> Union[float, np.ndarray]:
16
+ """
17
+ Planck function
18
+ Reference: (below 5.40)
19
+ """
20
+ return 2 * h_erg_s * nu_sm1**3 / c_cm_sm1**2 / (exp(h_erg_s * nu_sm1 / kB_erg_Km1 / T_K) - 1)
21
+
22
+
23
+ def nu_larmor(magnetic_field_gauss: np.ndarray) -> np.ndarray:
24
+ """
25
+ Larmor frequency in Hz
26
+ Reference: (3.10)
27
+ """
28
+ return magnetic_field_gauss * mu0_erg_gaussm1 / h_erg_s
29
+
30
+
31
+ def energy_cmm1_to_frequency_hz(energy_cmm1: Union[float, np.ndarray]) -> Union[float, np.ndarray]:
32
+ return c_cm_sm1 * energy_cmm1
33
+
34
+
35
+ def lambda_cm_to_frequency_hz(lambda_cm: Union[float, np.ndarray]) -> Union[float, np.ndarray]:
36
+ return c_cm_sm1 / lambda_cm
37
+
38
+
39
+ def frequency_hz_to_lambda_A(frequency_hz: Union[float, np.ndarray]) -> Union[float, np.ndarray]:
40
+ return c_cm_sm1 / frequency_hz * 1e8
41
+
42
+
43
+ def energy_cmm1_to_erg(energy_cmm1: float) -> float:
44
+ return h_erg_s * c_cm_sm1 * energy_cmm1
@@ -0,0 +1,120 @@
1
+ """
2
+ TODO
3
+ TODO This file needs improved documentation.
4
+ TODO
5
+ """
6
+
7
+ from functools import lru_cache
8
+
9
+ import numpy as np
10
+ import sympy
11
+ from numpy import cos, exp, sin
12
+ from sympy.physics.wigner import wigner_d
13
+
14
+ from solrat.common.constants import sqrt2, sqrt3
15
+ from solrat.engine.functions.general import delta, m1p
16
+ from solrat.engine.functions.looping import PROJECTION, fromto
17
+ from solrat.engine.generators.nested_loops import nested_loops
18
+
19
+
20
+ class WignerD:
21
+ """
22
+ Wigner D function.
23
+ alpha, beta, gamma are Euler angles in radians.
24
+ Typically, we have alpha = chi, beta = theta, gamma = gamma (see Fig. 5.14).
25
+ """
26
+
27
+ def __init__(self, alpha, beta, gamma, K_max):
28
+ self.d = {}
29
+ for K in fromto(0, K_max):
30
+ # Note: sympy uses a different convention for angles, so I perform under-the-hood conversion here.
31
+ self.d[K] = wigner_d(J=int(K), alpha=-alpha, beta=-beta, gamma=-gamma)
32
+
33
+ @lru_cache(maxsize=None)
34
+ def __call__(self, K, P, Q):
35
+ result = np.array(sympy.N(self.d[K][int(K - P), int(K - Q)])).astype(np.complex128)
36
+ return result
37
+
38
+
39
+ def t_K_P(K, P, stokes_component_index):
40
+ """
41
+ t{K, P}(i)
42
+ Reference: Table 5.5
43
+ This is implemented primarily to validate Wigner D functions
44
+ """
45
+ if K == 0:
46
+ return delta(P, 0) * delta(stokes_component_index, 0)
47
+ if K == 1:
48
+ return sqrt3 / sqrt2 * delta(P, 0) * delta(stokes_component_index, 3)
49
+ return (
50
+ 1 / sqrt2 * delta(P, 0) * delta(stokes_component_index, 0)
51
+ - sqrt3 / 2 * (delta(P, -2) + delta(P, 2)) * delta(stokes_component_index, 1)
52
+ + 1j * sqrt3 / 2 * (delta(P, -2) - delta(P, 2)) * delta(stokes_component_index, 2)
53
+ )
54
+
55
+
56
+ # @lru_cache(maxsize=None)
57
+ def T_K_Q_double_rotation(K, Q, stokes_component_index, D_inverse_omega: WignerD, D_magnetic: WignerD):
58
+ """
59
+ (5.159), (2.74), (5.122)
60
+ Two consecutive D rotations within T tensor.
61
+ """
62
+ result = 0
63
+ for P, Qʹ in nested_loops(P=PROJECTION(K), Qʹ=PROJECTION(K)):
64
+ result = result + t_K_P(
65
+ K=K,
66
+ P=P,
67
+ stokes_component_index=stokes_component_index,
68
+ ) * D_inverse_omega(
69
+ K=K, P=P, Q=Qʹ
70
+ ) * D_magnetic(K=K, P=Qʹ, Q=Q)
71
+
72
+ return result
73
+
74
+
75
+ @lru_cache(maxsize=None)
76
+ def T_K_Q(K, Q, stokes_component_index, chi, theta, gamma):
77
+ """
78
+ T{K, Q}(i, Omega)
79
+
80
+ Reference: Table 5.6
81
+ See also: (5.159), (5.160)
82
+ """
83
+ if Q < 0:
84
+ return m1p(Q) * T_K_Q(K, -Q, stokes_component_index, chi, theta, gamma).conjugate()
85
+
86
+ if stokes_component_index == 0:
87
+ if K == 0:
88
+ return 1 + 0j
89
+ if K == 1:
90
+ return 0 + 0j
91
+ if Q == 0:
92
+ return 0.5 / sqrt2 * (3 * cos(theta) ** 2 - 1) + 0j
93
+ if Q == 1:
94
+ return -0.5 * sqrt3 * sin(theta) * cos(theta) * exp(1j * chi)
95
+ return 0.25 * sqrt3 * sin(theta) ** 2 * exp(2j * chi)
96
+ if stokes_component_index == 1:
97
+ if K <= 1:
98
+ return 0 + 0j
99
+ if Q == 0:
100
+ return -1.5 / sqrt2 * cos(2 * gamma) * sin(theta) ** 2 + 0j
101
+ if Q == 1:
102
+ return -0.5 * sqrt3 * (cos(2 * gamma) * cos(theta) + 1j * sin(2 * gamma)) * sin(theta) * exp(1j * chi)
103
+ return (
104
+ -0.25 * sqrt3 * (cos(2 * gamma) * (1 + cos(theta) ** 2) + 2j * sin(2 * gamma) * cos(theta)) * exp(2j * chi)
105
+ )
106
+ if stokes_component_index == 2:
107
+ if K <= 1:
108
+ return 0 + 0j
109
+ if Q == 0:
110
+ return 1.5 / sqrt2 * sin(2 * gamma) * sin(theta) ** 2 + 0j
111
+ if Q == 1:
112
+ return 0.5 * sqrt3 * (sin(2 * gamma) * cos(theta) - 1j * cos(2 * gamma)) * sin(theta) * exp(1j * chi)
113
+ return (
114
+ 0.25 * sqrt3 * (sin(2 * gamma) * (1 + cos(theta) ** 2) - 2j * cos(2 * gamma) * cos(theta)) * exp(2j * chi)
115
+ )
116
+ if K == 0 or K == 2:
117
+ return 0 + 0j
118
+ if Q == 0:
119
+ return sqrt3 / sqrt2 * cos(theta) + 0j
120
+ return -0.5 * sqrt3 * sin(theta) * exp(1j * chi)
@@ -0,0 +1,45 @@
1
+ """
2
+ TODO
3
+ TODO This file needs improved documentation.
4
+ TODO
5
+ """
6
+
7
+ import numpy as np
8
+ from numpy import exp
9
+
10
+
11
+ def _voigt(nu: np.ndarray, a: float) -> np.ndarray:
12
+ """
13
+ complex Voigt profile at relative frequency nu with damping factor a.
14
+ Reference: Humlíček, J. (1982). JQSRT, doi:10.1016/0022-4073(82)90078-4
15
+
16
+ Note: HAZEL2 for some reason uses an expansion in terms of the Dawson's integral for a < 1e-3, see
17
+ https://github.com/aasensio/hazel2/blob/master/src/hazel/maths.f90#L1076.
18
+ They motivate it by some numerical instabilities near a->0.
19
+
20
+ I didn't find any such issues here in SolRaT, and the Humlíček expansion approach matches the results from
21
+ Dawson's integral expansion very closely, so I decided to use the Humlíček expansion for all a.
22
+ """
23
+
24
+ t = a - 1j * nu
25
+ s = abs(nu) + a
26
+ u = t * t
27
+
28
+ if s >= 15:
29
+ return t * 0.5641896 / (0.5 + u)
30
+ if s >= 5.5:
31
+ return t * (1.410474 + u * 0.5641896) / (0.75 + u * (3 + u))
32
+ if a >= 0.195 * abs(nu) - 0.176:
33
+ return (16.4955 + t * (20.20933 + t * (11.96482 + t * (3.778987 + t * 0.5642236)))) / (
34
+ 16.4955 + t * (38.82363 + t * (39.27121 + t * (21.69274 + t * (6.699398 + t))))
35
+ )
36
+ w4 = t * (
37
+ 36183.31 - u * (3321.9905 - u * (1540.787 - u * (219.0313 - u * (35.76683 - u * (1.320522 - u * 0.56419)))))
38
+ )
39
+ v4 = 32066.6 - u * (
40
+ 24322.84 - u * (9022.228 - u * (2186.181 - u * (364.2191 - u * (61.57037 - u * (1.841439 - u)))))
41
+ )
42
+ return exp(u) - w4 / v4
43
+
44
+
45
+ voigt = np.vectorize(_voigt)
@@ -0,0 +1,216 @@
1
+ """
2
+ TODO
3
+ TODO This file needs improved documentation.
4
+ TODO
5
+ """
6
+
7
+ from functools import lru_cache
8
+
9
+ import numpy as np
10
+ from numpy import abs, sqrt
11
+
12
+ from solrat.engine.functions.general import fact2
13
+
14
+
15
+ def _w3j_doubled_argument(j1_doubled, j2_doubled, j3_doubled, m1_doubled, m2_doubled, m3_doubled):
16
+ """
17
+ float Wigner 3J symbol where all arguments are doubled to be integer
18
+
19
+ j1_doubled = 2 * J1: int
20
+ j2_doubled = 2 * J2: int
21
+ j3_doubled = 2 * J3: int
22
+ m1_doubled = 2 * M1: int
23
+ m2_doubled = 2 * M2: int
24
+ m3_doubled = 2 * M3: int
25
+
26
+ Reference: Appendix A1
27
+
28
+ ( J1 J2 J3 )
29
+ ( M1 M2 M3 )
30
+ """
31
+
32
+ assert m1_doubled + m2_doubled + m3_doubled == 0, "M1 + M2 + M3 != 0."
33
+
34
+ if (abs(m1_doubled) > j1_doubled) or (abs(m2_doubled) > j2_doubled) or (abs(m3_doubled) > j3_doubled):
35
+ return 0.0
36
+ a = j1_doubled + j2_doubled
37
+ if j3_doubled > a:
38
+ return 0.0
39
+ b = j1_doubled - j2_doubled
40
+ if j3_doubled < abs(b):
41
+ return 0.0
42
+ j_sum = j3_doubled + a
43
+ c = j1_doubled - m1_doubled
44
+ d = j2_doubled - m2_doubled
45
+
46
+ assert j_sum % 2 == 0, "J1 + J2 + J3 != even."
47
+ assert c % 2 == 0, "J1 - M1 != even."
48
+ assert d % 2 == 0, "J2 - M2 != even."
49
+
50
+ e = j3_doubled - j2_doubled + m1_doubled
51
+ f = j3_doubled - j1_doubled - m2_doubled
52
+ z_min = max(0, -e, -f)
53
+ g = a - j3_doubled
54
+ h = j2_doubled + m2_doubled
55
+ z_max = min(g, h, c)
56
+ result = 0.0
57
+ for z in range(int(z_min), int(z_max) + 1, 2):
58
+ denominator = fact2(z) * fact2(g - z) * fact2(c - z) * fact2(h - z) * fact2(e + z) * fact2(f + z)
59
+ if z % 4 != 0:
60
+ denominator = -denominator
61
+ result += 1 / denominator
62
+ cc1 = fact2(g) * fact2(j3_doubled + b) * fact2(j3_doubled - b) / fact2(j_sum + 2)
63
+ cc2 = (
64
+ fact2(j1_doubled + m1_doubled)
65
+ * fact2(c)
66
+ * fact2(h)
67
+ * fact2(d)
68
+ * fact2(j3_doubled - m3_doubled)
69
+ * fact2(j3_doubled + m3_doubled)
70
+ )
71
+ result *= sqrt(cc1 * cc2)
72
+ if (b - m3_doubled) % 4 != 0:
73
+ result = -result
74
+ return result
75
+
76
+
77
+ def _w6j_doubled_argument(j1_doubled, j2_doubled, j3_doubled, l1_doubled, l2_doubled, l3_doubled):
78
+ """
79
+ float Wigner 6J symbol where all arguments are doubled to be integer
80
+
81
+ j1_doubled = 2 * J1: int
82
+ j2_doubled = 2 * J2: int
83
+ j3_doubled = 2 * J3: int
84
+ l1_doubled = 2 * L1: int
85
+ l2_doubled = 2 * L2: int
86
+ l3_doubled = 2 * L3: int
87
+
88
+ Reference: Appendix A1
89
+
90
+ { J1 J2 J3 }
91
+ { L1 L2 L3 }
92
+ """
93
+ a = j1_doubled + j2_doubled
94
+ b = j1_doubled - j2_doubled
95
+ c = j1_doubled + l2_doubled
96
+ d = j1_doubled - l2_doubled
97
+ e = l1_doubled + j2_doubled
98
+ f = l1_doubled - j2_doubled
99
+ g = l1_doubled + l2_doubled
100
+ h = l1_doubled - l2_doubled
101
+
102
+ if (a < j3_doubled) or (c < l3_doubled) or (e < l3_doubled) or (g < j3_doubled):
103
+ # logging.warning("Performance warning: J1 + J2 < J3 or L1 + L2 < L3 or L1 + J2 < L3 or L1 + L2 < J3")
104
+ return 0.0
105
+
106
+ if (abs(b) > j3_doubled) or (abs(d) > l3_doubled) or (abs(f) > l3_doubled) or (abs(h) > j3_doubled):
107
+ # logging.warning("Performance warning: J1 - J2 > J3 or J1 - L2 > L3 or L1 - J2 > L3 or L1 - L2 > J3")
108
+ return 0.0
109
+
110
+ sum_1 = a + j3_doubled
111
+ sum_2 = c + l3_doubled
112
+ sum_3 = e + l3_doubled
113
+ sum_4 = g + j3_doubled
114
+
115
+ assert sum_1 % 2 == 0, "J1 + J2 + J3 != even."
116
+ assert sum_2 % 2 == 0, "J1 + L2 + L3 != even."
117
+ assert sum_3 % 2 == 0, "L1 + J2 + L3 != even."
118
+
119
+ w_min = max(sum_1, sum_2, sum_3, sum_4)
120
+ i = a + g
121
+ j = j2_doubled + j3_doubled + l2_doubled + l3_doubled
122
+ k = j3_doubled + j1_doubled + l3_doubled + l1_doubled
123
+ w_max = min(i, j, k)
124
+
125
+ result = 0.0
126
+ for w in range(int(w_min), int(w_max) + 1, 2):
127
+ denominator = (
128
+ fact2(w - sum_1)
129
+ * fact2(w - sum_2)
130
+ * fact2(w - sum_3)
131
+ * fact2(w - sum_4)
132
+ * fact2(i - w)
133
+ * fact2(j - w)
134
+ * fact2(k - w)
135
+ )
136
+ if w % 4 != 0:
137
+ denominator = -denominator
138
+ result += fact2(w + 2) / denominator
139
+
140
+ theta1 = fact2(a - j3_doubled) * fact2(j3_doubled + b) * fact2(j3_doubled - b) / fact2(sum_1 + 2)
141
+ theta2 = fact2(c - l3_doubled) * fact2(l3_doubled + d) * fact2(l3_doubled - d) / fact2(sum_2 + 2)
142
+ theta3 = fact2(e - l3_doubled) * fact2(l3_doubled + f) * fact2(l3_doubled - f) / fact2(sum_3 + 2)
143
+ theta4 = fact2(g - j3_doubled) * fact2(j3_doubled + h) * fact2(j3_doubled - h) / fact2(sum_4 + 2)
144
+ result = result * sqrt(theta1 * theta2 * theta3 * theta4)
145
+ return result
146
+
147
+
148
+ def _w9j_doubled_argument(
149
+ j1_doubled,
150
+ j2_doubled,
151
+ j3_doubled,
152
+ j4_doubled,
153
+ j5_doubled,
154
+ j6_doubled,
155
+ j7_doubled,
156
+ j8_doubled,
157
+ j9_doubled,
158
+ ):
159
+ """
160
+ float Wigner 9J symbol where all arguments are doubled to be integer
161
+
162
+ j1_doubled = 2 * J1: int
163
+ j2_doubled = 2 * J2: int
164
+ j3_doubled = 2 * J3: int
165
+ j4_doubled = 2 * J4: int
166
+ j5_doubled = 2 * J5: int
167
+ j6_doubled = 2 * J6: int
168
+ j7_doubled = 2 * J7: int
169
+ j8_doubled = 2 * J8: int
170
+ j9_doubled = 2 * J9: int
171
+
172
+ Reference: Appendix A1
173
+
174
+ { J1 J2 J3 }
175
+ { J4 J5 J6 }
176
+ { J7 J8 J9 }
177
+ """
178
+ k_min = max(
179
+ abs(j1_doubled - j9_doubled),
180
+ abs(j4_doubled - j8_doubled),
181
+ abs(j2_doubled - j6_doubled),
182
+ )
183
+
184
+ k_max = min(
185
+ abs(j1_doubled + j9_doubled),
186
+ abs(j4_doubled + j8_doubled),
187
+ abs(j2_doubled + j6_doubled),
188
+ )
189
+ result = 0
190
+ for k in range(int(k_min), int(k_max) + 1, 2):
191
+ s = -1 if k % 2 != 0 else 1
192
+ x1 = _w6j_doubled_argument(j1_doubled, j9_doubled, k, j8_doubled, j4_doubled, j7_doubled)
193
+ x2 = _w6j_doubled_argument(j2_doubled, j6_doubled, k, j4_doubled, j8_doubled, j5_doubled)
194
+ x3 = _w6j_doubled_argument(j1_doubled, j9_doubled, k, j6_doubled, j2_doubled, j3_doubled)
195
+ result += s * x1 * x2 * x3 * (k + 1)
196
+ return result
197
+
198
+
199
+ # vectorize
200
+ _w3j_doubled_argument_vec = np.vectorize(_w3j_doubled_argument)
201
+ _w6j_doubled_argument_vec = np.vectorize(_w6j_doubled_argument)
202
+
203
+
204
+ # @lru_cache(maxsize=None)
205
+ def wigner_3j(j1, j2, j3, m1, m2, m3):
206
+ return _w3j_doubled_argument_vec(j1 * 2, j2 * 2, j3 * 2, m1 * 2, m2 * 2, m3 * 2)
207
+
208
+
209
+ # @lru_cache(maxsize=None)
210
+ def wigner_6j(j1, j2, j3, l1, l2, l3):
211
+ return _w6j_doubled_argument_vec(j1 * 2, j2 * 2, j3 * 2, l1 * 2, l2 * 2, l3 * 2)
212
+
213
+
214
+ @lru_cache(maxsize=None)
215
+ def wigner_9j(j1, j2, j3, j4, j5, j6, j7, j8, j9):
216
+ return _w9j_doubled_argument(j1 * 2, j2 * 2, j3 * 2, j4 * 2, j5 * 2, j6 * 2, j7 * 2, j8 * 2, j9 * 2)
File without changes
File without changes