diffpes 2026.3.1__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.
- diffpes/__init__.py +65 -0
- diffpes/inout/__init__.py +88 -0
- diffpes/inout/chgcar.py +459 -0
- diffpes/inout/doscar.py +248 -0
- diffpes/inout/eigenval.py +258 -0
- diffpes/inout/hdf5.py +818 -0
- diffpes/inout/helpers.py +295 -0
- diffpes/inout/kpoints.py +433 -0
- diffpes/inout/plotting.py +757 -0
- diffpes/inout/poscar.py +174 -0
- diffpes/inout/procar.py +331 -0
- diffpes/inout/py.typed +0 -0
- diffpes/maths/__init__.py +68 -0
- diffpes/maths/dipole.py +368 -0
- diffpes/maths/gaunt.py +496 -0
- diffpes/maths/spherical_harmonics.py +336 -0
- diffpes/py.typed +0 -0
- diffpes/radial/__init__.py +47 -0
- diffpes/radial/bessel.py +198 -0
- diffpes/radial/integrate.py +128 -0
- diffpes/radial/wavefunctions.py +335 -0
- diffpes/simul/__init__.py +146 -0
- diffpes/simul/broadening.py +258 -0
- diffpes/simul/crosssections.py +196 -0
- diffpes/simul/expanded.py +1021 -0
- diffpes/simul/forward.py +586 -0
- diffpes/simul/oam.py +112 -0
- diffpes/simul/polarization.py +443 -0
- diffpes/simul/py.typed +0 -0
- diffpes/simul/resolution.py +122 -0
- diffpes/simul/self_energy.py +142 -0
- diffpes/simul/spectrum.py +1116 -0
- diffpes/simul/workflow.py +424 -0
- diffpes/tightb/__init__.py +68 -0
- diffpes/tightb/diagonalize.py +340 -0
- diffpes/tightb/hamiltonian.py +286 -0
- diffpes/tightb/projections.py +134 -0
- diffpes/types/__init__.py +162 -0
- diffpes/types/aliases.py +52 -0
- diffpes/types/bands.py +1064 -0
- diffpes/types/dos.py +510 -0
- diffpes/types/geometry.py +306 -0
- diffpes/types/kpath.py +365 -0
- diffpes/types/params.py +496 -0
- diffpes/types/py.typed +0 -0
- diffpes/types/radial_params.py +482 -0
- diffpes/types/self_energy.py +271 -0
- diffpes/types/tb_model.py +531 -0
- diffpes/types/volumetric.py +608 -0
- diffpes/utils/__init__.py +30 -0
- diffpes/utils/math.py +255 -0
- diffpes/utils/py.typed +0 -0
- diffpes-2026.3.1.dist-info/METADATA +176 -0
- diffpes-2026.3.1.dist-info/RECORD +55 -0
- diffpes-2026.3.1.dist-info/WHEEL +4 -0
diffpes/utils/math.py
ADDED
|
@@ -0,0 +1,255 @@
|
|
|
1
|
+
"""Mathematical utility functions for ARPES simulations.
|
|
2
|
+
|
|
3
|
+
Extended Summary
|
|
4
|
+
----------------
|
|
5
|
+
Provides JAX-compatible implementations of the Faddeeva function
|
|
6
|
+
(complex error function) and data normalization routines used
|
|
7
|
+
throughout the ARPES simulation pipeline.
|
|
8
|
+
|
|
9
|
+
Routine Listings
|
|
10
|
+
----------------
|
|
11
|
+
:func:`faddeeva`
|
|
12
|
+
Faddeeva function via Taylor series expansion.
|
|
13
|
+
:func:`zscore_normalize`
|
|
14
|
+
Z-score (zero-mean, unit-variance) normalization.
|
|
15
|
+
|
|
16
|
+
Notes
|
|
17
|
+
-----
|
|
18
|
+
The Faddeeva implementation uses a 64-term Taylor series
|
|
19
|
+
derived from the ODE w'(z) = -2z w(z) + 2i/sqrt(pi),
|
|
20
|
+
giving double-precision accuracy for |z| < 6.
|
|
21
|
+
"""
|
|
22
|
+
|
|
23
|
+
import math
|
|
24
|
+
|
|
25
|
+
import jax
|
|
26
|
+
import jax.numpy as jnp
|
|
27
|
+
from beartype import beartype
|
|
28
|
+
from jaxtyping import Array, Complex, Float, Int, jaxtyped
|
|
29
|
+
|
|
30
|
+
_N_TAYLOR: int = 64
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
def _faddeeva_taylor_coeffs() -> Complex[Array, " N"]:
|
|
34
|
+
r"""Taylor coefficients of w(z) = exp(-z^2) erfc(-iz) via JAX scan.
|
|
35
|
+
|
|
36
|
+
Extended Summary
|
|
37
|
+
----------------
|
|
38
|
+
Computes the first ``_N_TAYLOR`` coefficients of the Taylor expansion
|
|
39
|
+
of the Faddeeva function about the origin.
|
|
40
|
+
|
|
41
|
+
**Mathematical derivation:**
|
|
42
|
+
|
|
43
|
+
The Faddeeva function satisfies the ODE:
|
|
44
|
+
|
|
45
|
+
.. math::
|
|
46
|
+
|
|
47
|
+
w'(z) = -2z \, w(z) + \frac{2i}{\sqrt{\pi}}
|
|
48
|
+
|
|
49
|
+
Substituting the power series ansatz :math:`w(z) = \sum_{n=0}^\infty
|
|
50
|
+
a_n z^n` and matching coefficients of :math:`z^n` on both sides
|
|
51
|
+
yields the two-term recurrence:
|
|
52
|
+
|
|
53
|
+
.. math::
|
|
54
|
+
|
|
55
|
+
a_0 = 1, \quad a_1 = \frac{2i}{\sqrt{\pi}}, \quad
|
|
56
|
+
a_{n+1} = \frac{-2 \, a_{n-1}}{n+1} \quad (n \ge 1)
|
|
57
|
+
|
|
58
|
+
Note that the recurrence couples even-indexed coefficients among
|
|
59
|
+
themselves and odd-indexed coefficients among themselves (the
|
|
60
|
+
series separates into even and odd parts). Even coefficients are
|
|
61
|
+
real and odd coefficients are purely imaginary, reflecting the
|
|
62
|
+
symmetry :math:`w(-z) = 2 e^{-z^2} - w(z)`.
|
|
63
|
+
|
|
64
|
+
**Implementation:**
|
|
65
|
+
|
|
66
|
+
The recurrence is implemented with ``jax.lax.scan`` (no Python
|
|
67
|
+
for-loop) for efficiency and XLA compatibility. The scan carry
|
|
68
|
+
is the pair :math:`(a_{n-1}, a_n)`, and at each step n (0-indexed),
|
|
69
|
+
the next coefficient is :math:`a_{n+2} = -2 a_n / (n + 2)`. The
|
|
70
|
+
scan produces coefficients :math:`a_2` through :math:`a_{N-1}`,
|
|
71
|
+
which are concatenated with the seed values :math:`a_0, a_1` to
|
|
72
|
+
form the complete coefficient vector.
|
|
73
|
+
|
|
74
|
+
The resulting array is reversed (descending power order) at module
|
|
75
|
+
level and stored in ``_W_POLY`` for use with ``jnp.polyval``
|
|
76
|
+
(Horner's method).
|
|
77
|
+
|
|
78
|
+
Returns
|
|
79
|
+
-------
|
|
80
|
+
coeffs : Complex[Array, " N"]
|
|
81
|
+
Taylor coefficients :math:`[a_0, a_1, \ldots, a_{N-1}]` in
|
|
82
|
+
ascending power order, where ``N = _N_TAYLOR``.
|
|
83
|
+
"""
|
|
84
|
+
c0: Complex[Array, " "] = jnp.array(1.0 + 0j, dtype=jnp.complex128)
|
|
85
|
+
c1: Complex[Array, " "] = jnp.array(
|
|
86
|
+
2.0j / math.sqrt(math.pi), dtype=jnp.complex128
|
|
87
|
+
)
|
|
88
|
+
|
|
89
|
+
def body(
|
|
90
|
+
carry: tuple[Complex[Array, " "], Complex[Array, " "]],
|
|
91
|
+
n: Int[Array, ""],
|
|
92
|
+
) -> tuple[
|
|
93
|
+
tuple[Complex[Array, " "], Complex[Array, " "]],
|
|
94
|
+
Complex[Array, " "],
|
|
95
|
+
]:
|
|
96
|
+
c_prev, c_curr = carry
|
|
97
|
+
next_c: Complex[Array, " "] = (-2.0 * c_prev) / (
|
|
98
|
+
jnp.asarray(n, dtype=jnp.float64) + 2.0
|
|
99
|
+
)
|
|
100
|
+
return (c_curr, next_c), next_c
|
|
101
|
+
|
|
102
|
+
rest: Complex[Array, " N2"]
|
|
103
|
+
_, rest = jax.lax.scan(
|
|
104
|
+
body,
|
|
105
|
+
(c0, c1),
|
|
106
|
+
jnp.arange(_N_TAYLOR - 2, dtype=jnp.int32),
|
|
107
|
+
)
|
|
108
|
+
full: Complex[Array, " N"] = jnp.concatenate([c0[None], c1[None], rest])
|
|
109
|
+
return full
|
|
110
|
+
|
|
111
|
+
|
|
112
|
+
_W_POLY: Complex[Array, " N"] = _faddeeva_taylor_coeffs()[::-1]
|
|
113
|
+
|
|
114
|
+
|
|
115
|
+
@jaxtyped(typechecker=beartype)
|
|
116
|
+
def faddeeva(
|
|
117
|
+
z: Complex[Array, " ..."],
|
|
118
|
+
) -> Complex[Array, " ..."]:
|
|
119
|
+
r"""Evaluate the Faddeeva function w(z) = exp(-z^2) erfc(-iz).
|
|
120
|
+
|
|
121
|
+
Extended Summary
|
|
122
|
+
----------------
|
|
123
|
+
Computes the Faddeeva (scaled complex complementary error) function
|
|
124
|
+
for arbitrary complex arrays using a precomputed Taylor polynomial
|
|
125
|
+
evaluated via Horner's method.
|
|
126
|
+
|
|
127
|
+
The Faddeeva function is defined as:
|
|
128
|
+
|
|
129
|
+
.. math::
|
|
130
|
+
|
|
131
|
+
w(z) = e^{-z^2} \operatorname{erfc}(-iz)
|
|
132
|
+
= e^{-z^2} \left(1 + \frac{2i}{\sqrt{\pi}}
|
|
133
|
+
\int_0^z e^{t^2} dt \right)
|
|
134
|
+
|
|
135
|
+
It is closely related to the Voigt profile used in spectroscopy:
|
|
136
|
+
the Voigt function is the real part of the Faddeeva function
|
|
137
|
+
evaluated along the imaginary axis. In ARPES simulations, it
|
|
138
|
+
appears when convolving Lorentzian lifetime broadening with
|
|
139
|
+
Gaussian instrumental resolution.
|
|
140
|
+
|
|
141
|
+
**Implementation:**
|
|
142
|
+
|
|
143
|
+
1. **Cast to complex128** -- convert the input to ``jnp.complex128``
|
|
144
|
+
via ``jnp.asarray`` to ensure sufficient precision for the 64-term
|
|
145
|
+
polynomial evaluation.
|
|
146
|
+
|
|
147
|
+
2. **Evaluate via Horner's method** -- call ``jnp.polyval`` with
|
|
148
|
+
the module-level constant ``_W_POLY`` (coefficients stored in
|
|
149
|
+
descending-power order, i.e. reversed from the natural
|
|
150
|
+
:math:`a_0, \ldots, a_{N-1}` ordering) and the cast input. The
|
|
151
|
+
``unroll=8`` hint allows XLA to pipeline the inner loop for
|
|
152
|
+
better throughput on accelerators.
|
|
153
|
+
|
|
154
|
+
The Taylor coefficients are derived from the ODE
|
|
155
|
+
:math:`w'(z) = -2z w(z) + 2i/\sqrt{\pi}` and precomputed at
|
|
156
|
+
module import time by `_faddeeva_taylor_coeffs`. See that
|
|
157
|
+
function's docstring for the full recurrence derivation.
|
|
158
|
+
|
|
159
|
+
Parameters
|
|
160
|
+
----------
|
|
161
|
+
z : Complex[Array, " ..."]
|
|
162
|
+
Complex argument(s), arbitrary shape.
|
|
163
|
+
|
|
164
|
+
Returns
|
|
165
|
+
-------
|
|
166
|
+
w : Complex[Array, " ..."]
|
|
167
|
+
Faddeeva function values, same shape as ``z``.
|
|
168
|
+
|
|
169
|
+
Notes
|
|
170
|
+
-----
|
|
171
|
+
Accuracy is approximately double-precision (~15 significant digits)
|
|
172
|
+
for |z| < 6. For |z| >= 6 the Taylor series converges slowly and
|
|
173
|
+
an asymptotic expansion or continued-fraction representation should
|
|
174
|
+
be used instead. The current implementation does **not** fall back
|
|
175
|
+
to an asymptotic form, so callers must ensure inputs stay within the
|
|
176
|
+
convergence domain.
|
|
177
|
+
|
|
178
|
+
The function is fully differentiable via JAX autodiff, which
|
|
179
|
+
is important for gradient-based optimization of broadening
|
|
180
|
+
parameters in inverse-fitting workflows.
|
|
181
|
+
"""
|
|
182
|
+
z_c: Complex[Array, " ..."] = jnp.asarray(z, dtype=jnp.complex128)
|
|
183
|
+
w: Complex[Array, " ..."] = jnp.polyval(_W_POLY, z_c, unroll=8)
|
|
184
|
+
return w
|
|
185
|
+
|
|
186
|
+
|
|
187
|
+
@jaxtyped(typechecker=beartype)
|
|
188
|
+
def zscore_normalize(
|
|
189
|
+
data: Float[Array, " ..."],
|
|
190
|
+
) -> Float[Array, " ..."]:
|
|
191
|
+
r"""Apply z-score normalization (zero-mean, unit-variance).
|
|
192
|
+
|
|
193
|
+
Extended Summary
|
|
194
|
+
----------------
|
|
195
|
+
Transforms an arbitrary float array so that the output has zero mean
|
|
196
|
+
and unit standard deviation, which is a common preprocessing step
|
|
197
|
+
before comparing simulated and experimental ARPES spectra. The
|
|
198
|
+
z-score transformation is:
|
|
199
|
+
|
|
200
|
+
.. math::
|
|
201
|
+
|
|
202
|
+
\hat{x}_i = \frac{x_i - \bar{x}}{\sigma}
|
|
203
|
+
|
|
204
|
+
where :math:`\bar{x}` is the global mean and :math:`\sigma` is
|
|
205
|
+
the population standard deviation (:math:`\text{ddof}=0`).
|
|
206
|
+
|
|
207
|
+
**Implementation details:**
|
|
208
|
+
|
|
209
|
+
1. **Compute statistics** -- calculate the global mean and standard
|
|
210
|
+
deviation of the input array using ``jnp.mean`` and ``jnp.std``
|
|
211
|
+
(population std, i.e. ddof=0).
|
|
212
|
+
|
|
213
|
+
2. **Guard against zero std** -- if the standard deviation is
|
|
214
|
+
exactly zero (constant array), replace it with 1.0 via
|
|
215
|
+
``jnp.where(std > 0, std, 1.0)`` to avoid division-by-zero.
|
|
216
|
+
This produces an all-zeros output for constant inputs, which is
|
|
217
|
+
the mathematically sensible limit. The ``jnp.where`` formulation
|
|
218
|
+
is gradient-safe: JAX will not propagate gradients through the
|
|
219
|
+
zero-std branch.
|
|
220
|
+
|
|
221
|
+
3. **Normalize** -- compute ``(data - mean) / safe_std`` element-wise
|
|
222
|
+
and return.
|
|
223
|
+
|
|
224
|
+
Parameters
|
|
225
|
+
----------
|
|
226
|
+
data : Float[Array, " ..."]
|
|
227
|
+
Input data array of any shape.
|
|
228
|
+
|
|
229
|
+
Returns
|
|
230
|
+
-------
|
|
231
|
+
normalized : Float[Array, " ..."]
|
|
232
|
+
Normalized data with mean 0 and standard deviation 1
|
|
233
|
+
(or all zeros if the input is constant).
|
|
234
|
+
|
|
235
|
+
Notes
|
|
236
|
+
-----
|
|
237
|
+
The normalization is computed over **all** elements (global mean
|
|
238
|
+
and std), not per-axis. For per-axis normalization, reshape the
|
|
239
|
+
data and call this function on each slice separately.
|
|
240
|
+
|
|
241
|
+
This function is differentiable via JAX autodiff with respect to
|
|
242
|
+
the input ``data``. The gradient propagates through both the
|
|
243
|
+
mean-subtraction and the division by standard deviation.
|
|
244
|
+
"""
|
|
245
|
+
mean_val: Float[Array, " "] = jnp.mean(data)
|
|
246
|
+
std_val: Float[Array, " "] = jnp.std(data)
|
|
247
|
+
safe_std: Float[Array, " "] = jnp.where(std_val > 0.0, std_val, 1.0)
|
|
248
|
+
normalized: Float[Array, " ..."] = (data - mean_val) / safe_std
|
|
249
|
+
return normalized
|
|
250
|
+
|
|
251
|
+
|
|
252
|
+
__all__: list[str] = [
|
|
253
|
+
"faddeeva",
|
|
254
|
+
"zscore_normalize",
|
|
255
|
+
]
|
diffpes/utils/py.typed
ADDED
|
File without changes
|
|
@@ -0,0 +1,176 @@
|
|
|
1
|
+
Metadata-Version: 2.3
|
|
2
|
+
Name: diffpes
|
|
3
|
+
Version: 2026.3.1
|
|
4
|
+
Summary: Differentiable ARPES simulations in JAX
|
|
5
|
+
Keywords: ARPES,MBE,JAX,Differentiable Programming
|
|
6
|
+
Author: Debangshu Mukherjee, Jacob Cook
|
|
7
|
+
Author-email: Debangshu Mukherjee <mukherjeed@ornl.gov>, Jacob Cook <cookjl2@ornl.gov>
|
|
8
|
+
License: MIT License
|
|
9
|
+
|
|
10
|
+
Copyright (c) 2026 Debangshu Mukherjee, Oak Ridge National Laboratory
|
|
11
|
+
|
|
12
|
+
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:
|
|
13
|
+
|
|
14
|
+
The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
|
|
15
|
+
|
|
16
|
+
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 THE AUTHORS OR COPYRIGHT HOLDERS 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.
|
|
17
|
+
Classifier: Development Status :: 4 - Beta
|
|
18
|
+
Classifier: Intended Audience :: Science/Research
|
|
19
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
20
|
+
Classifier: Operating System :: OS Independent
|
|
21
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
22
|
+
Classifier: Programming Language :: Python :: 3.14
|
|
23
|
+
Classifier: Topic :: Scientific/Engineering :: Physics
|
|
24
|
+
Classifier: Topic :: Scientific/Engineering :: Chemistry
|
|
25
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
26
|
+
Classifier: Typing :: Typed
|
|
27
|
+
Requires-Dist: numpy>=2.2.1
|
|
28
|
+
Requires-Dist: matplotlib>=3.10.0
|
|
29
|
+
Requires-Dist: jax>=0.7.0 ; sys_platform == 'win64'
|
|
30
|
+
Requires-Dist: jax>=0.7.0 ; sys_platform == 'darwin'
|
|
31
|
+
Requires-Dist: jax>=0.7.0 ; sys_platform == 'linux'
|
|
32
|
+
Requires-Dist: jaxtyping>=0.3.0
|
|
33
|
+
Requires-Dist: chex>=0.1.85
|
|
34
|
+
Requires-Dist: beartype>=0.21.0
|
|
35
|
+
Requires-Dist: jupyter>=1.1.1
|
|
36
|
+
Requires-Dist: ipykernel>=7.2.0
|
|
37
|
+
Requires-Dist: h5py>=3.15.1
|
|
38
|
+
Requires-Dist: diffpes[docs,test,dev,notebooks,cuda] ; extra == 'all'
|
|
39
|
+
Requires-Dist: jax[cuda12]>=0.7.0 ; platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'cuda'
|
|
40
|
+
Requires-Dist: nvidia-cudnn-cu12>=9.5.0 ; platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'cuda'
|
|
41
|
+
Requires-Dist: diffpes[docs,test,notebooks] ; extra == 'dev'
|
|
42
|
+
Requires-Dist: black[jupyter]>=25.1.0 ; extra == 'dev'
|
|
43
|
+
Requires-Dist: build>=1.2.2.post1 ; extra == 'dev'
|
|
44
|
+
Requires-Dist: twine>=6.1.0 ; extra == 'dev'
|
|
45
|
+
Requires-Dist: ruff>=0.12.9 ; extra == 'dev'
|
|
46
|
+
Requires-Dist: pygount>=3.1.0 ; extra == 'dev'
|
|
47
|
+
Requires-Dist: pre-commit>=4.5.1 ; extra == 'dev'
|
|
48
|
+
Requires-Dist: isort>=6.0.1 ; extra == 'dev'
|
|
49
|
+
Requires-Dist: ty>=0.0.19 ; extra == 'dev'
|
|
50
|
+
Requires-Dist: black>=26.1.0 ; extra == 'dev'
|
|
51
|
+
Requires-Dist: jupyter-black>=0.4.0 ; extra == 'dev'
|
|
52
|
+
Requires-Dist: diffpes[dev,cuda] ; extra == 'dev-cuda'
|
|
53
|
+
Requires-Dist: sphinx>=7.0.0 ; extra == 'docs'
|
|
54
|
+
Requires-Dist: sphinx-rtd-theme>=3.0.2 ; extra == 'docs'
|
|
55
|
+
Requires-Dist: nbsphinx>=0.9.7 ; extra == 'docs'
|
|
56
|
+
Requires-Dist: myst-parser>=2.0.0 ; extra == 'docs'
|
|
57
|
+
Requires-Dist: ipykernel>=6.29.5 ; extra == 'docs'
|
|
58
|
+
Requires-Dist: nbconvert>=7.16.6 ; extra == 'docs'
|
|
59
|
+
Requires-Dist: furo>=2025.7.19 ; extra == 'docs'
|
|
60
|
+
Requires-Dist: sphinx-autodoc-typehints>=3.0.1 ; extra == 'docs'
|
|
61
|
+
Requires-Dist: numpydoc>=1.9.0 ; extra == 'docs'
|
|
62
|
+
Requires-Dist: pydoclint>=0.7.3 ; extra == 'docs'
|
|
63
|
+
Requires-Dist: interrogate>=1.7.0 ; extra == 'docs'
|
|
64
|
+
Requires-Dist: ipywidgets>=8.1.0 ; extra == 'notebooks'
|
|
65
|
+
Requires-Dist: ipykernel>=6.29.5 ; extra == 'notebooks'
|
|
66
|
+
Requires-Dist: nbconvert>=7.16.6 ; extra == 'notebooks'
|
|
67
|
+
Requires-Dist: pytest>=8.3.5 ; extra == 'test'
|
|
68
|
+
Requires-Dist: pytest-cov>=6.0.0 ; extra == 'test'
|
|
69
|
+
Requires-Dist: pytest-xdist>=3.0.0 ; extra == 'test'
|
|
70
|
+
Requires-Dist: chex>=0.1.89 ; extra == 'test'
|
|
71
|
+
Maintainer: Debangshu Mukherjee
|
|
72
|
+
Maintainer-email: Debangshu Mukherjee <mukherjeed@ornl.gov>
|
|
73
|
+
Requires-Python: >=3.12, <3.15
|
|
74
|
+
Provides-Extra: all
|
|
75
|
+
Provides-Extra: cuda
|
|
76
|
+
Provides-Extra: dev
|
|
77
|
+
Provides-Extra: dev-cuda
|
|
78
|
+
Provides-Extra: docs
|
|
79
|
+
Provides-Extra: notebooks
|
|
80
|
+
Provides-Extra: test
|
|
81
|
+
Description-Content-Type: text/markdown
|
|
82
|
+
|
|
83
|
+
# diffpes
|
|
84
|
+
|
|
85
|
+
JAX-based ARPES simulation toolkit with Python-native APIs.
|
|
86
|
+
|
|
87
|
+
## Expanded-input workflows
|
|
88
|
+
|
|
89
|
+
The package includes expanded-input wrappers that let you call the
|
|
90
|
+
simulator with plain arrays/scalars while still running JAX kernels.
|
|
91
|
+
|
|
92
|
+
### Function mapping
|
|
93
|
+
|
|
94
|
+
- `ARPES_simulation_Novice` -> `diffpes.simul.simulate_novice_expanded`
|
|
95
|
+
- `ARPES_simulation_Basic` -> `diffpes.simul.simulate_basic_expanded`
|
|
96
|
+
- `ARPES_simulation_Basicplus` -> `diffpes.simul.simulate_basicplus_expanded`
|
|
97
|
+
- `ARPES_simulation_Advanced` -> `diffpes.simul.simulate_advanced_expanded`
|
|
98
|
+
- `ARPES_simulation_Expert` -> `diffpes.simul.simulate_expert_expanded`
|
|
99
|
+
- `ARPES_simulation_SOC` -> `diffpes.simul.simulate_soc_expanded`
|
|
100
|
+
- Dynamic dispatch by level -> `diffpes.simul.simulate_expanded`
|
|
101
|
+
(use `level="soc"` with `surface_spin` for SOC)
|
|
102
|
+
|
|
103
|
+
### Notes
|
|
104
|
+
|
|
105
|
+
- Default energy-axis padding behavior:
|
|
106
|
+
`min(eigenbands)-1` to `max(eigenbands)+1`.
|
|
107
|
+
- Incident angles for expanded wrappers are interpreted in degrees.
|
|
108
|
+
- Wrappers return the standard `ArpesSpectrum` PyTree.
|
|
109
|
+
|
|
110
|
+
### Python indexing conventions
|
|
111
|
+
|
|
112
|
+
Use standard Python/NumPy indexing everywhere (zero-based, end-exclusive).
|
|
113
|
+
|
|
114
|
+
- Non-s orbitals: `slice(1, 9)` -> indices 1..8
|
|
115
|
+
- p orbitals: `slice(1, 4)` -> indices 1..3
|
|
116
|
+
- d orbitals: `slice(4, 9)` -> indices 4..8
|
|
117
|
+
|
|
118
|
+
Do not use MATLAB-style indexing notation in Python code.
|
|
119
|
+
|
|
120
|
+
### Example
|
|
121
|
+
|
|
122
|
+
```python
|
|
123
|
+
import jax.numpy as jnp
|
|
124
|
+
|
|
125
|
+
from diffpes.simul import simulate_expanded
|
|
126
|
+
|
|
127
|
+
# [nkpt, nband]
|
|
128
|
+
eigenbands = jnp.linspace(-2.0, 0.5, 100).reshape(20, 5)
|
|
129
|
+
# [nkpt, nband, natom, 9]
|
|
130
|
+
surface_orb = jnp.ones((20, 5, 2, 9)) * 0.1
|
|
131
|
+
|
|
132
|
+
spectrum = simulate_expanded(
|
|
133
|
+
level="advanced",
|
|
134
|
+
eigenbands=eigenbands,
|
|
135
|
+
surface_orb=surface_orb,
|
|
136
|
+
ef=0.0,
|
|
137
|
+
sigma=0.04,
|
|
138
|
+
fidelity=2500,
|
|
139
|
+
temperature=15.0,
|
|
140
|
+
photon_energy=11.0,
|
|
141
|
+
polarization="unpolarized",
|
|
142
|
+
incident_theta=45.0,
|
|
143
|
+
incident_phi=0.0,
|
|
144
|
+
polarization_angle=0.0,
|
|
145
|
+
)
|
|
146
|
+
```
|
|
147
|
+
|
|
148
|
+
## Test coverage
|
|
149
|
+
|
|
150
|
+
Test coverage measures which lines of source code are executed
|
|
151
|
+
during tests. Run it with:
|
|
152
|
+
|
|
153
|
+
```bash
|
|
154
|
+
source .venv/bin/activate
|
|
155
|
+
pytest tests/ --cov=src/diffpes --cov-report=term-missing
|
|
156
|
+
```
|
|
157
|
+
|
|
158
|
+
To get as close to 100% as possible:
|
|
159
|
+
|
|
160
|
+
1. **Simul and types** — Already well covered. Any new branch
|
|
161
|
+
(e.g. new polarization or dispatch level) should have a
|
|
162
|
+
corresponding test.
|
|
163
|
+
2. **Expanded dispatch** — Test every `simulate_expanded(level=...)`
|
|
164
|
+
branch (novice, basic, basicplus, advanced, expert, soc) and
|
|
165
|
+
the unknown-level `ValueError`.
|
|
166
|
+
3. **HDF5** — Round-trip all PyTree types; test error paths
|
|
167
|
+
(unknown type on load, missing group, unsupported type on save).
|
|
168
|
+
4. **VASP file readers** (`read_doscar`, `read_eigenval`, `read_kpoints`,
|
|
169
|
+
`read_poscar`, `read_procar`) — Add tests that call each reader
|
|
170
|
+
on minimal in-repo fixture files (e.g. under `tests/fixtures/`)
|
|
171
|
+
so the parsing code paths are executed.
|
|
172
|
+
5. **Plotting** — Exercise the public plotting API in tests (or
|
|
173
|
+
accept lower coverage for GUI-oriented code).
|
|
174
|
+
6. **Edge branches** — Cover optional arguments (e.g.
|
|
175
|
+
`make_band_structure(..., kpoint_weights=...)`) and error
|
|
176
|
+
messages so one-off branches are hit.
|
|
@@ -0,0 +1,55 @@
|
|
|
1
|
+
diffpes/__init__.py,sha256=dyrIAP5eT5itDHSrqHN5qgFjIWYIpCC0M0YHkRd8w50,1781
|
|
2
|
+
diffpes/inout/__init__.py,sha256=eJ3ZXLESDUQWHsB_M0-yVzCRbeOdXASk52YUc_LICoE,2565
|
|
3
|
+
diffpes/inout/chgcar.py,sha256=O3JXU0d-1YxHx86LPBYQAs8IVk0FmeDR0mv5bpD_cRI,15865
|
|
4
|
+
diffpes/inout/doscar.py,sha256=ITcMVcNio2Y1XTsXSLBqwwyuUOv2oCAjEEZh63nMnFo,9685
|
|
5
|
+
diffpes/inout/eigenval.py,sha256=kibxceLicm_5HlO0DduD21d2rIcJwzoIUZ1yFHlTLIE,9743
|
|
6
|
+
diffpes/inout/hdf5.py,sha256=jbnNdW8C9m45MKfi4p5W9UWWdb6aEHieTUGePfuqxWw,24880
|
|
7
|
+
diffpes/inout/helpers.py,sha256=Q8wHM2_bXXTRCJeO3rZ3FnhE7S5qkys_nFWbX8ILPpc,10250
|
|
8
|
+
diffpes/inout/kpoints.py,sha256=x4PrTpjTLet7r9ZIT5DwS7KEFe0U23mCHpxMp3wSiPc,13944
|
|
9
|
+
diffpes/inout/plotting.py,sha256=tLEHMqE6KJqWdSXXB6Cl2I4BoGQ5xIGMAMt099uYXI0,22759
|
|
10
|
+
diffpes/inout/poscar.py,sha256=o6BUd0A3_V6jboES7jCGbkUtFLfWveKMasdueXMVh5s,6937
|
|
11
|
+
diffpes/inout/procar.py,sha256=wH1_nMFFdE5384yhYTgKpZuDi8OVBrKQ9usZcbZgrKM,12203
|
|
12
|
+
diffpes/inout/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
13
|
+
diffpes/maths/__init__.py,sha256=rrxgTCe7dT-Yk2prtuEN0qmDPvrzMEQVJYpSsRJFctc,2406
|
|
14
|
+
diffpes/maths/dipole.py,sha256=wbmK7d65akLB-ibK6QEE6WT32PCstQDnomisaF-azKk,13039
|
|
15
|
+
diffpes/maths/gaunt.py,sha256=T0xRsKwROpKIW2xfkcH8OpuZsxrFxYjKKxGyGsEJB00,16668
|
|
16
|
+
diffpes/maths/spherical_harmonics.py,sha256=YJ3Lg3lLkHjPnMC2OLk0OLxgue9PCoycUWV2Eq96RSE,10282
|
|
17
|
+
diffpes/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
18
|
+
diffpes/radial/__init__.py,sha256=GVo0oUbeivBXNCRwdK5-psQqi1uh8JoDjSBWzZoJzaw,1435
|
|
19
|
+
diffpes/radial/bessel.py,sha256=3p0AvYVCg1mUM0-zeco2TZHMYxGGBJu3LHC35JxhmKI,6012
|
|
20
|
+
diffpes/radial/integrate.py,sha256=SeSiUGpSD2zJRsChQYmgczgUzx84NqRr_NmGl8yKXu8,4250
|
|
21
|
+
diffpes/radial/wavefunctions.py,sha256=tz3IVyD3wVw22G3NJOhn2HsoQ2_rN8_BkEqyB5T9i7I,10433
|
|
22
|
+
diffpes/simul/__init__.py,sha256=oH1dK4_hlRHgACOgPxmWHJaD1Ao8DTZg8PaUxH52png,4653
|
|
23
|
+
diffpes/simul/broadening.py,sha256=krkuYfnGtP3UrhCCJrSByvK7ECg19byekeU4STMNz5k,8881
|
|
24
|
+
diffpes/simul/crosssections.py,sha256=Hg4ytjOjGu9KMgc0JtFUC8fPK-luiM2n8BW9jm2edBY,6831
|
|
25
|
+
diffpes/simul/expanded.py,sha256=GiMzNht2zDmAePqVqg_mzMYx7eeiAyEPWtL2QSYTQzE,36000
|
|
26
|
+
diffpes/simul/forward.py,sha256=ZzJu_8FVUfULSPvPTE6bnM1tjXxkHjjRA5s-zMlkNk4,22901
|
|
27
|
+
diffpes/simul/oam.py,sha256=YUAZ4-ey69kOlsxYvCsogcyVv2CisKGki8z_qEBeOO0,3737
|
|
28
|
+
diffpes/simul/polarization.py,sha256=aEgjdqxbtwldwc90YhKm15Kv8zAVlcMZDzF81nF1iqk,14614
|
|
29
|
+
diffpes/simul/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
30
|
+
diffpes/simul/resolution.py,sha256=JPRNpCr0FZ7AOvgaU8Fk0LNrLuxTOcqhqT1NUiv70G8,4772
|
|
31
|
+
diffpes/simul/self_energy.py,sha256=BlG138FqnRc1tgrswH7Opd-0GdexUq4as2Cyi99kfSs,5375
|
|
32
|
+
diffpes/simul/spectrum.py,sha256=dH2P8Dts1oUfAV8F6O0q_KnMxM_WMH0QUEXVp0qRjqY,42380
|
|
33
|
+
diffpes/simul/workflow.py,sha256=uTi6MXXJA46ZI1hJuDneUr9CrN0v_KR1a37jiJMuyKk,14132
|
|
34
|
+
diffpes/tightb/__init__.py,sha256=kZwk-QAimOeT8eqeH3zQ0bkg6e31dsAY8q_0pcxy-EU,2261
|
|
35
|
+
diffpes/tightb/diagonalize.py,sha256=NF6WOfybPm71QGd4Rza837id4lf9vrMoUdramRqXMBs,12420
|
|
36
|
+
diffpes/tightb/hamiltonian.py,sha256=Hb-XzcdDuIbDBPZRATpqL_VT58VCr5ojO1H2_8CezDg,9679
|
|
37
|
+
diffpes/tightb/projections.py,sha256=WMrWnI_XQm4lk5RgzG_Wo_g2uuG-7_2HS7joGQ3EgaY,4477
|
|
38
|
+
diffpes/types/__init__.py,sha256=90jceGqMY3D9UtYATXN-OcQt4wjJn0dSM0gnkaoFdy8,4368
|
|
39
|
+
diffpes/types/aliases.py,sha256=99JfYDr1P4xuEHvrBgvcwtl8hN5lHle1Hk_UUCdhxPA,1553
|
|
40
|
+
diffpes/types/bands.py,sha256=kACWB0jaG5dAnB6mtBi-v3uCQL4Wo1YrgXUVM2JK4dI,40857
|
|
41
|
+
diffpes/types/dos.py,sha256=MVPTPx9uIoYfVax5YwCtwlREmkk93PnicHns4S2dSiM,18763
|
|
42
|
+
diffpes/types/geometry.py,sha256=kPrpRopD7OJzki4IAiel2hEpPcmm5jAQUW4NVQ8m3LM,11019
|
|
43
|
+
diffpes/types/kpath.py,sha256=x6VgtKc-eYZm3kUJFjktLkJuXrQcuEbPpJYO1DfNB8g,13032
|
|
44
|
+
diffpes/types/params.py,sha256=5xVFsD4rkvu70M5-vu2JpkjlUpkWlPNZM8RggOUCYuA,17712
|
|
45
|
+
diffpes/types/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
46
|
+
diffpes/types/radial_params.py,sha256=pdFanc7GLoi7BWuq-sxUxujgqcoVwZeiYfk7q9gHbY4,17874
|
|
47
|
+
diffpes/types/self_energy.py,sha256=zucn_V9PYQSLd_ZEqjOA0RhBg7LWKkJkrh2Xm8R63YY,10287
|
|
48
|
+
diffpes/types/tb_model.py,sha256=ZriOaGX2LnixXYzJLXC_EMOMpZ28E0EQ2pMZNjm0dNY,19706
|
|
49
|
+
diffpes/types/volumetric.py,sha256=j9IMyY2PMkCqhXv4gI5PEEHYLbQrAdVSLNdwGrLKlPw,23506
|
|
50
|
+
diffpes/utils/__init__.py,sha256=HJj0aTrTiaZYZQ8qXoGG2fE0O81CrHynz19XVzT3AC4,960
|
|
51
|
+
diffpes/utils/math.py,sha256=SGC42JSZm4wKuDWBL-JLU__VOViAvP-5E8rODl1HZr8,8927
|
|
52
|
+
diffpes/utils/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
53
|
+
diffpes-2026.3.1.dist-info/WHEEL,sha256=Uo4e6VmJM8J_cLBTtiWBRVQkc1yUEkw94xaL0B_lH6c,80
|
|
54
|
+
diffpes-2026.3.1.dist-info/METADATA,sha256=B6KdLnSSMRB2cYghm-096eoTUHjaxkzD5PM_RPgzFog,7631
|
|
55
|
+
diffpes-2026.3.1.dist-info/RECORD,,
|