spectral_flux 0.1.3__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.
- spectral_flux-0.1.3/LICENSE +21 -0
- spectral_flux-0.1.3/PKG-INFO +27 -0
- spectral_flux-0.1.3/README.md +2 -0
- spectral_flux-0.1.3/pyproject.toml +112 -0
- spectral_flux-0.1.3/pyproject.toml.orig +115 -0
- spectral_flux-0.1.3/src/spectral_flux/__init__.py +1 -0
- spectral_flux-0.1.3/src/spectral_flux/advection.py +89 -0
- spectral_flux-0.1.3/src/spectral_flux/coarse_calc.py +69 -0
- spectral_flux-0.1.3/src/spectral_flux/fluidsf_compat.py +78 -0
- spectral_flux-0.1.3/src/spectral_flux/fourier_calc.py +183 -0
- spectral_flux-0.1.3/src/spectral_flux/sf_au.py +285 -0
- spectral_flux-0.1.3/src/spectral_flux/sf_int.py +190 -0
- spectral_flux-0.1.3/src/spectral_flux/sf_ln.py +267 -0
- spectral_flux-0.1.3/src/spectral_flux/utils.py +241 -0
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Peter Carlip
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1,27 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: spectral_flux
|
|
3
|
+
Version: 0.1.3
|
|
4
|
+
Summary: Python tools for calculating spectral flux from fluid velocities
|
|
5
|
+
Author: pcarlip
|
|
6
|
+
Author-email: pcarlip <peter.carlip@proton.me>
|
|
7
|
+
License-Expression: MIT
|
|
8
|
+
License-File: LICENSE
|
|
9
|
+
Requires-Dist: cupy-xarray>=0.1.5
|
|
10
|
+
Requires-Dist: numpy>=2.5.2
|
|
11
|
+
Requires-Dist: pandas>=3.0.5
|
|
12
|
+
Requires-Dist: scipy>=1.18.1
|
|
13
|
+
Requires-Dist: xarray>=2026.7.0
|
|
14
|
+
Requires-Dist: xarray-extras>=0.6.0
|
|
15
|
+
Requires-Dist: xrft
|
|
16
|
+
Requires-Dist: xrscipy>=3.1.0
|
|
17
|
+
Requires-Dist: cupy-cuda12x ; extra == 'cupy12'
|
|
18
|
+
Requires-Dist: cupy-cuda13x ; extra == 'cupy13'
|
|
19
|
+
Requires-Python: >=3.13
|
|
20
|
+
Project-URL: Homepage, https://github.com/pcarlip/spec_flux
|
|
21
|
+
Project-URL: Issues, https://github.com/pcarlip/spec_flux/issues
|
|
22
|
+
Provides-Extra: cupy12
|
|
23
|
+
Provides-Extra: cupy13
|
|
24
|
+
Description-Content-Type: text/markdown
|
|
25
|
+
|
|
26
|
+
# spec_flux
|
|
27
|
+
Python tools for calculating spectral flux from fluid velocities
|
|
@@ -0,0 +1,112 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "spectral_flux"
|
|
3
|
+
version = "0.1.3"
|
|
4
|
+
description = "Python tools for calculating spectral flux from fluid velocities"
|
|
5
|
+
readme = "README.md"
|
|
6
|
+
requires-python = ">=3.13"
|
|
7
|
+
dependencies = [
|
|
8
|
+
"cupy-xarray>=0.1.5",
|
|
9
|
+
"numpy>=2.5.2",
|
|
10
|
+
"pandas>=3.0.5",
|
|
11
|
+
"scipy>=1.18.1",
|
|
12
|
+
"xarray>=2026.7.0",
|
|
13
|
+
"xarray-extras>=0.6.0",
|
|
14
|
+
"xrft",
|
|
15
|
+
"xrscipy>=3.1.0",
|
|
16
|
+
]
|
|
17
|
+
license = "MIT"
|
|
18
|
+
license-files = ["LICENSE"]
|
|
19
|
+
|
|
20
|
+
[[project.authors]]
|
|
21
|
+
name = "pcarlip"
|
|
22
|
+
email = "peter.carlip@proton.me"
|
|
23
|
+
|
|
24
|
+
[project.urls]
|
|
25
|
+
Homepage = "https://github.com/pcarlip/spec_flux"
|
|
26
|
+
Issues = "https://github.com/pcarlip/spec_flux/issues"
|
|
27
|
+
|
|
28
|
+
[project.optional-dependencies]
|
|
29
|
+
cupy12 = ["cupy-cuda12x"]
|
|
30
|
+
cupy13 = ["cupy-cuda13x"]
|
|
31
|
+
|
|
32
|
+
[project.scripts]
|
|
33
|
+
spec-flux = "spec_flux:main"
|
|
34
|
+
|
|
35
|
+
[tool.uv]
|
|
36
|
+
conflicts = [[
|
|
37
|
+
{ extra = "cuda12" },
|
|
38
|
+
{ extra = "cuda13" },
|
|
39
|
+
]]
|
|
40
|
+
|
|
41
|
+
[tool.uv.sources.xrft]
|
|
42
|
+
git = "https://github.com/pcarlip/xrft.git"
|
|
43
|
+
branch = "cupy"
|
|
44
|
+
|
|
45
|
+
[tool.ruff]
|
|
46
|
+
exclude = [
|
|
47
|
+
".git",
|
|
48
|
+
".ipynb_checkpoints",
|
|
49
|
+
".mypy_cache",
|
|
50
|
+
".venv",
|
|
51
|
+
".vscode",
|
|
52
|
+
"__pypackages__",
|
|
53
|
+
"venv",
|
|
54
|
+
]
|
|
55
|
+
line-length = 90
|
|
56
|
+
indent-width = 4
|
|
57
|
+
target-version = "py313"
|
|
58
|
+
|
|
59
|
+
[tool.ruff.lint]
|
|
60
|
+
select = [
|
|
61
|
+
"F",
|
|
62
|
+
"B",
|
|
63
|
+
"UP",
|
|
64
|
+
"I",
|
|
65
|
+
"ANN",
|
|
66
|
+
"E",
|
|
67
|
+
"A",
|
|
68
|
+
"C4",
|
|
69
|
+
"ICN",
|
|
70
|
+
"PIE",
|
|
71
|
+
"SIM",
|
|
72
|
+
"PLE",
|
|
73
|
+
"PLC",
|
|
74
|
+
"PLR",
|
|
75
|
+
"PLW",
|
|
76
|
+
"NPY",
|
|
77
|
+
"RUF",
|
|
78
|
+
"PD",
|
|
79
|
+
"PYI",
|
|
80
|
+
"PERF",
|
|
81
|
+
"C",
|
|
82
|
+
"FURB",
|
|
83
|
+
]
|
|
84
|
+
ignore = [
|
|
85
|
+
"F401",
|
|
86
|
+
"PLR0913",
|
|
87
|
+
"E203",
|
|
88
|
+
"B018",
|
|
89
|
+
"PLR0917",
|
|
90
|
+
"PLR0914",
|
|
91
|
+
"RUF003",
|
|
92
|
+
"PLR0915",
|
|
93
|
+
]
|
|
94
|
+
fixable = ["ALL"]
|
|
95
|
+
unfixable = []
|
|
96
|
+
|
|
97
|
+
[tool.ruff.format]
|
|
98
|
+
quote-style = "double"
|
|
99
|
+
indent-style = "space"
|
|
100
|
+
skip-magic-trailing-comma = false
|
|
101
|
+
line-ending = "auto"
|
|
102
|
+
|
|
103
|
+
[build-system]
|
|
104
|
+
requires = ["uv_build>=0.12.7,<0.13.0"]
|
|
105
|
+
build-backend = "uv_build"
|
|
106
|
+
|
|
107
|
+
[dependency-groups]
|
|
108
|
+
dev = ["ruff>=0.16.5"]
|
|
109
|
+
test = [
|
|
110
|
+
"fluidsf>=0.2.2",
|
|
111
|
+
"pytest>=9.1.1",
|
|
112
|
+
]
|
|
@@ -0,0 +1,115 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "spectral_flux"
|
|
3
|
+
version = "0.1.3"
|
|
4
|
+
description = "Python tools for calculating spectral flux from fluid velocities"
|
|
5
|
+
readme = "README.md"
|
|
6
|
+
authors = [{ name = "pcarlip", email = "peter.carlip@proton.me" }]
|
|
7
|
+
requires-python = ">=3.13"
|
|
8
|
+
dependencies = [
|
|
9
|
+
"cupy-xarray>=0.1.5",
|
|
10
|
+
"numpy>=2.5.2",
|
|
11
|
+
"pandas>=3.0.5",
|
|
12
|
+
"scipy>=1.18.1",
|
|
13
|
+
"xarray>=2026.7.0",
|
|
14
|
+
"xarray-extras>=0.6.0",
|
|
15
|
+
"xrft",
|
|
16
|
+
"xrscipy>=3.1.0",
|
|
17
|
+
]
|
|
18
|
+
license = "MIT"
|
|
19
|
+
license-files = ["LICENSE"]
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
[project.urls]
|
|
23
|
+
Homepage = "https://github.com/pcarlip/spec_flux"
|
|
24
|
+
Issues = "https://github.com/pcarlip/spec_flux/issues"
|
|
25
|
+
|
|
26
|
+
[project.optional-dependencies]
|
|
27
|
+
cupy12 = ["cupy-cuda12x"]
|
|
28
|
+
cupy13 = ["cupy-cuda13x"]
|
|
29
|
+
|
|
30
|
+
[tool.uv]
|
|
31
|
+
conflicts = [[{ extra = "cuda12" }, { extra = "cuda13" }]]
|
|
32
|
+
|
|
33
|
+
[project.scripts]
|
|
34
|
+
spec-flux = "spec_flux:main"
|
|
35
|
+
|
|
36
|
+
[build-system]
|
|
37
|
+
requires = ["uv_build>=0.12.7,<0.13.0"]
|
|
38
|
+
build-backend = "uv_build"
|
|
39
|
+
|
|
40
|
+
[tool.uv.sources]
|
|
41
|
+
xrft = { git = "https://github.com/pcarlip/xrft.git", branch = "cupy" }
|
|
42
|
+
# currently using my own branch of xrft which supports cupy
|
|
43
|
+
# will switch to main branch if my PR is accepted
|
|
44
|
+
|
|
45
|
+
[dependency-groups]
|
|
46
|
+
dev = ["ruff>=0.16.5"]
|
|
47
|
+
test = ["fluidsf>=0.2.2", "pytest>=9.1.1"]
|
|
48
|
+
|
|
49
|
+
[tool.ruff]
|
|
50
|
+
exclude = [
|
|
51
|
+
".git",
|
|
52
|
+
".ipynb_checkpoints",
|
|
53
|
+
".mypy_cache",
|
|
54
|
+
".venv",
|
|
55
|
+
".vscode",
|
|
56
|
+
"__pypackages__",
|
|
57
|
+
"venv",
|
|
58
|
+
]
|
|
59
|
+
|
|
60
|
+
line-length = 90
|
|
61
|
+
indent-width = 4
|
|
62
|
+
target-version = "py313"
|
|
63
|
+
|
|
64
|
+
[tool.ruff.lint]
|
|
65
|
+
select = [
|
|
66
|
+
"F",
|
|
67
|
+
"B",
|
|
68
|
+
"UP",
|
|
69
|
+
"I",
|
|
70
|
+
"ANN",
|
|
71
|
+
"E",
|
|
72
|
+
"A",
|
|
73
|
+
"C4",
|
|
74
|
+
"ICN",
|
|
75
|
+
"PIE",
|
|
76
|
+
"SIM",
|
|
77
|
+
"PLE",
|
|
78
|
+
"PLC",
|
|
79
|
+
"PLR",
|
|
80
|
+
"PLW",
|
|
81
|
+
"NPY",
|
|
82
|
+
"RUF",
|
|
83
|
+
"PD",
|
|
84
|
+
"PYI",
|
|
85
|
+
"PERF",
|
|
86
|
+
"C",
|
|
87
|
+
"FURB",
|
|
88
|
+
]
|
|
89
|
+
|
|
90
|
+
ignore = [
|
|
91
|
+
"F401",
|
|
92
|
+
"PLR0913",
|
|
93
|
+
"E203",
|
|
94
|
+
"B018",
|
|
95
|
+
"PLR0917",
|
|
96
|
+
"PLR0914",
|
|
97
|
+
"RUF003",
|
|
98
|
+
"PLR0915",
|
|
99
|
+
]
|
|
100
|
+
fixable = ["ALL"]
|
|
101
|
+
unfixable = []
|
|
102
|
+
|
|
103
|
+
|
|
104
|
+
[tool.ruff.format]
|
|
105
|
+
# Like Black, use double quotes for strings.
|
|
106
|
+
quote-style = "double"
|
|
107
|
+
|
|
108
|
+
# Like Black, indent with spaces, rather than tabs.
|
|
109
|
+
indent-style = "space"
|
|
110
|
+
|
|
111
|
+
# Like Black, respect magic trailing commas.
|
|
112
|
+
skip-magic-trailing-comma = false
|
|
113
|
+
|
|
114
|
+
# Like Black, automatically detect the appropriate line ending.
|
|
115
|
+
line-ending = "auto"
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
|
|
@@ -0,0 +1,89 @@
|
|
|
1
|
+
from typing import Literal
|
|
2
|
+
|
|
3
|
+
import cupy_xarray
|
|
4
|
+
import xarray as xr
|
|
5
|
+
|
|
6
|
+
from .utils import (
|
|
7
|
+
Axis,
|
|
8
|
+
GradMethod,
|
|
9
|
+
krange_fft,
|
|
10
|
+
ndarray,
|
|
11
|
+
xp_fft,
|
|
12
|
+
)
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
def spectral_der(vel_hat: ndarray, k_grid: ndarray) -> ndarray:
|
|
16
|
+
"""Approximate dv/dx as IFT(ik FT(v)), with cupy acceleration
|
|
17
|
+
|
|
18
|
+
Parameters
|
|
19
|
+
----------
|
|
20
|
+
vel_hat : np.ndarray | cp.ndarray
|
|
21
|
+
Fourier transform of a velocity component
|
|
22
|
+
k_grid : np.ndarray | cp.ndarray
|
|
23
|
+
meshgrid component of k values along an axis
|
|
24
|
+
|
|
25
|
+
Returns
|
|
26
|
+
-------
|
|
27
|
+
np.ndarray | cp.ndarray
|
|
28
|
+
Realspace derivative of velocity with respect to the direction from the k-grid
|
|
29
|
+
"""
|
|
30
|
+
xp, genfft = xp_fft(vel_hat)
|
|
31
|
+
return xp.real(genfft.ifftn(genfft.ifftshift(1j * k_grid * vel_hat)))
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
def advection_xr(
|
|
35
|
+
data: xr.Dataset,
|
|
36
|
+
axis: Axis,
|
|
37
|
+
method: GradMethod,
|
|
38
|
+
ax_names: tuple[str, str, str] = ("z_aac", "y_aca", "x_caa"),
|
|
39
|
+
edge_order: Literal[1, 2] = 1,
|
|
40
|
+
) -> xr.DataArray:
|
|
41
|
+
"""Calculate (realspace) advection of a velocity array
|
|
42
|
+
|
|
43
|
+
Parameters
|
|
44
|
+
----------
|
|
45
|
+
data : xr.Dataset
|
|
46
|
+
Dataset with velocity and (optionally) advection data
|
|
47
|
+
axis : Axis
|
|
48
|
+
Velocity direction to use
|
|
49
|
+
method : GradMethod
|
|
50
|
+
Method for calculating velocity gradients
|
|
51
|
+
ax_names : tuple[str, str, str], optional
|
|
52
|
+
names of z, y, and x axes in dataset, by default ("z_aac", "y_aca", "x_caa")
|
|
53
|
+
edge_order : Literal[1, 2], optional
|
|
54
|
+
Order of finite difference gradients if using GradMethod.numpy, by default 1
|
|
55
|
+
|
|
56
|
+
Returns
|
|
57
|
+
-------
|
|
58
|
+
xr.DataArray
|
|
59
|
+
xarray object containing advection along the specified axis
|
|
60
|
+
|
|
61
|
+
Raises
|
|
62
|
+
------
|
|
63
|
+
Exception
|
|
64
|
+
Attempt to use GradMethod.oceananigans (precalculated advection) when
|
|
65
|
+
input data doesn't include that information
|
|
66
|
+
"""
|
|
67
|
+
vel = (data["w"], data["v"], data["u"])[axis.value]
|
|
68
|
+
if method == GradMethod.precalc:
|
|
69
|
+
if "uadv" in data and "vadv" in data and "wadv" in data:
|
|
70
|
+
adv = (data["wadv"], data["vadv"], data["uadv"])[axis.value]
|
|
71
|
+
else:
|
|
72
|
+
raise Exception("Include advection arrays to use precalculated gradients")
|
|
73
|
+
elif method == GradMethod.numpy:
|
|
74
|
+
adv = (
|
|
75
|
+
data["u"] * vel.differentiate(ax_names[2], edge_order)
|
|
76
|
+
+ data["v"] * vel.differentiate(ax_names[1], edge_order)
|
|
77
|
+
+ data["w"] * vel.differentiate(ax_names[0], edge_order)
|
|
78
|
+
)
|
|
79
|
+
else:
|
|
80
|
+
xp, genfft = xp_fft(vel.data)
|
|
81
|
+
k_ranges = krange_fft(data)
|
|
82
|
+
k_mesh = xp.meshgrid(*k_ranges, indexing="ij")
|
|
83
|
+
vel_hat = genfft.fftshift(genfft.fftn(vel.data))
|
|
84
|
+
adv = (
|
|
85
|
+
data["w"] * spectral_der(vel_hat, k_mesh[0])
|
|
86
|
+
+ data["v"] * spectral_der(vel_hat, k_mesh[1])
|
|
87
|
+
+ data["u"] * spectral_der(vel_hat, k_mesh[2])
|
|
88
|
+
)
|
|
89
|
+
return adv
|
|
@@ -0,0 +1,69 @@
|
|
|
1
|
+
from collections.abc import Iterable
|
|
2
|
+
|
|
3
|
+
import cupy as cp
|
|
4
|
+
import cupy_xarray
|
|
5
|
+
import xarray as xr
|
|
6
|
+
from cupyx.scipy.ndimage import gaussian_filter
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
def pi_cg_gauss_nd(
|
|
10
|
+
data: xr.Dataset,
|
|
11
|
+
k: float,
|
|
12
|
+
skip_dims: tuple[str, ...] = ("time",),
|
|
13
|
+
vel_names: tuple[str, str, str] = ("u", "v", "w"),
|
|
14
|
+
axes: tuple[str, str, str] = ("z_aac", "y_aca", "x_caa"),
|
|
15
|
+
periodic: tuple[bool, ...] = (True, True, True),
|
|
16
|
+
) -> xr.DataArray:
|
|
17
|
+
|
|
18
|
+
# <f(s)> = ∫dr G(r)f(s+r), for which I use "gaussian_filter"
|
|
19
|
+
# τ_ij = <u_i u_j> - <u_i> <u_j>
|
|
20
|
+
# Π = -(∂_i <u_j>) τ_ij
|
|
21
|
+
|
|
22
|
+
smooth_dims = [i for i in axes if i not in skip_dims]
|
|
23
|
+
smooth_axes = [data[vel_names[0]].dims.index(i) for i in smooth_dims]
|
|
24
|
+
|
|
25
|
+
vels = tuple(data[i] for i in vel_names)
|
|
26
|
+
dz = float(data[axes[0]][1] - data[axes[0]][0])
|
|
27
|
+
dy = float(data[axes[1]][1] - data[axes[1]][0])
|
|
28
|
+
dx = float(data[axes[2]][1] - data[axes[2]][0])
|
|
29
|
+
size = ((1 / k) / dz, (1 / k) / dy, (1 / k) / dx)
|
|
30
|
+
|
|
31
|
+
modes = ["wrap" if i else "nearest" for i in periodic]
|
|
32
|
+
gauss_kwargs = {"sigma": size, "mode": modes, "axes": smooth_axes}
|
|
33
|
+
|
|
34
|
+
smoothed_vels = [
|
|
35
|
+
xr.apply_ufunc(gaussian_filter, vels[i], kwargs=gauss_kwargs) for i in range(3)
|
|
36
|
+
]
|
|
37
|
+
running_sum = xr.DataArray(0.0, {"time": data.time, "k": k})
|
|
38
|
+
|
|
39
|
+
for i in range(3):
|
|
40
|
+
for j in range(3):
|
|
41
|
+
tau_1 = xr.apply_ufunc(
|
|
42
|
+
gaussian_filter, vels[i] * vels[j], kwargs=gauss_kwargs
|
|
43
|
+
)
|
|
44
|
+
tau_2 = smoothed_vels[i] * smoothed_vels[j]
|
|
45
|
+
tau = tau_1 - tau_2
|
|
46
|
+
grad = vels[i].differentiate(axes[j], 2)
|
|
47
|
+
running_sum -= cp.mean(tau * grad).data.get()
|
|
48
|
+
|
|
49
|
+
return running_sum
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
def pi_cg_gauss_xr(
|
|
53
|
+
data: xr.Dataset,
|
|
54
|
+
k: float,
|
|
55
|
+
vel_names: tuple[str, str, str] = ("u", "v", "w"),
|
|
56
|
+
axes: tuple[str, str, str] = ("z_aac", "y_aca", "x_caa"),
|
|
57
|
+
periodic: tuple[bool, bool, bool] = (True, True, True),
|
|
58
|
+
) -> xr.DataArray:
|
|
59
|
+
return pi_cg_gauss_nd(data, k, vel_names=vel_names, axes=axes, periodic=periodic)
|
|
60
|
+
|
|
61
|
+
|
|
62
|
+
def pi_cg_lst_xr(data: xr.Dataset, k_cg: Iterable[float]) -> xr.DataArray:
|
|
63
|
+
return xr.concat([pi_cg_gauss_xr(data, k) for k in k_cg], "k")
|
|
64
|
+
|
|
65
|
+
|
|
66
|
+
def pi_cg_lst_nd(
|
|
67
|
+
data: xr.Dataset, k_cg: Iterable[float], skip_dims: tuple[str, ...] = ("time",)
|
|
68
|
+
) -> xr.DataArray:
|
|
69
|
+
return xr.concat([pi_cg_gauss_nd(data, k, skip_dims) for k in k_cg], "k")
|
|
@@ -0,0 +1,78 @@
|
|
|
1
|
+
import cupy as cp
|
|
2
|
+
import numpy as np
|
|
3
|
+
import xarray as xr
|
|
4
|
+
|
|
5
|
+
from .sf_au import sf_au_dir_xr
|
|
6
|
+
from .sf_ln import sf_ln_dir_xr
|
|
7
|
+
from .utils import Axis, ndarray
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
def fluidsf_compat(
|
|
11
|
+
u: ndarray,
|
|
12
|
+
v: ndarray,
|
|
13
|
+
w: ndarray,
|
|
14
|
+
x: ndarray,
|
|
15
|
+
y: ndarray,
|
|
16
|
+
z: ndarray,
|
|
17
|
+
sf_names: list[str],
|
|
18
|
+
) -> dict[str, np.ndarray]:
|
|
19
|
+
"""Get structure functions with compatible input ant output to fluidsf, using my
|
|
20
|
+
existing structure function methods
|
|
21
|
+
Arrays must be indexed as (z,y,x)
|
|
22
|
+
|
|
23
|
+
Parameters
|
|
24
|
+
----------
|
|
25
|
+
u : ndarray
|
|
26
|
+
velocity along x
|
|
27
|
+
v : ndarray
|
|
28
|
+
velocity along y
|
|
29
|
+
w : ndarray
|
|
30
|
+
velocity along z
|
|
31
|
+
x : ndarray
|
|
32
|
+
locations along x
|
|
33
|
+
y : ndarray
|
|
34
|
+
locations along y
|
|
35
|
+
z : ndarray
|
|
36
|
+
locations along z
|
|
37
|
+
sf_names : list[str]
|
|
38
|
+
structure functions to calculate (supports 'ASF_V', 'LL', and 'LLL')
|
|
39
|
+
|
|
40
|
+
Returns
|
|
41
|
+
-------
|
|
42
|
+
dict[str, np.ndarray]
|
|
43
|
+
dictionary of directional structure function names and values
|
|
44
|
+
"""
|
|
45
|
+
|
|
46
|
+
dims = ("z_aac", "y_aca", "x_caa")
|
|
47
|
+
ds = xr.Dataset(
|
|
48
|
+
{"u": (dims, u), "v": (dims, v), "w": (dims, w)},
|
|
49
|
+
coords={"z_aac": z, "y_aca": y, "x_caa": x},
|
|
50
|
+
)
|
|
51
|
+
|
|
52
|
+
L = len(z) // 2
|
|
53
|
+
M = len(y) // 2
|
|
54
|
+
N = len(x) // 2
|
|
55
|
+
|
|
56
|
+
dz = z[:L] - z[0]
|
|
57
|
+
dy = y[:M] - y[0]
|
|
58
|
+
dx = x[:N] - x[0]
|
|
59
|
+
|
|
60
|
+
out = {"x-diffs": dx, "y-diffs": dy, "z-diffs": dz}
|
|
61
|
+
if "ASF_V" in sf_names:
|
|
62
|
+
out["SF_advection_velocity_x"] = sf_au_dir_xr(ds, Axis.x).data
|
|
63
|
+
out["SF_advection_velocity_y"] = sf_au_dir_xr(ds, Axis.y).data
|
|
64
|
+
out["SF_advection_velocity_z"] = sf_au_dir_xr(ds, Axis.z).data
|
|
65
|
+
if "LLL" in sf_names:
|
|
66
|
+
out["SF_LLL_x"] = sf_ln_dir_xr(ds, Axis.x, 3).data
|
|
67
|
+
out["SF_LLL_y"] = sf_ln_dir_xr(ds, Axis.y, 3).data
|
|
68
|
+
out["SF_LLL_z"] = sf_ln_dir_xr(ds, Axis.z, 3).data
|
|
69
|
+
if "LL" in sf_names:
|
|
70
|
+
out["SF_LL_x"] = sf_ln_dir_xr(ds, Axis.x, 2).data
|
|
71
|
+
out["SF_LL_y"] = sf_ln_dir_xr(ds, Axis.y, 2).data
|
|
72
|
+
out["SF_LL_z"] = sf_ln_dir_xr(ds, Axis.z, 2).data
|
|
73
|
+
|
|
74
|
+
if cp.get_array_module(u).__name__ == "cupy":
|
|
75
|
+
for key, val in out.items():
|
|
76
|
+
out[key] = val.get()
|
|
77
|
+
|
|
78
|
+
return out
|
|
@@ -0,0 +1,183 @@
|
|
|
1
|
+
from collections.abc import Collection, Iterable
|
|
2
|
+
from typing import Literal
|
|
3
|
+
|
|
4
|
+
import cupy_xarray
|
|
5
|
+
import numpy as np
|
|
6
|
+
import xarray as xr
|
|
7
|
+
import xrft
|
|
8
|
+
from xrscipy.integrate import cumulative_simpson
|
|
9
|
+
|
|
10
|
+
from .advection import advection_xr
|
|
11
|
+
from .utils import Axis, GradMethod, axis_name
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
def pi_int_dir_xr(
|
|
15
|
+
data: xr.Dataset,
|
|
16
|
+
axis: Axis,
|
|
17
|
+
method: GradMethod,
|
|
18
|
+
edge_order: Literal[1, 2] = 2,
|
|
19
|
+
) -> xr.DataArray:
|
|
20
|
+
"""Generate a component of array to integrate: Re[FT(u)* • FT((u•∇)u)],
|
|
21
|
+
with optional cupy acceleration
|
|
22
|
+
|
|
23
|
+
Parameters
|
|
24
|
+
----------
|
|
25
|
+
data : xr.Dataset
|
|
26
|
+
Dataset with realspace velocity components (on the same grid)
|
|
27
|
+
axis : Axis
|
|
28
|
+
Axis along which to generate the component
|
|
29
|
+
method : GradMethod
|
|
30
|
+
Method for calculating gradients
|
|
31
|
+
edge_order : Literal[1, 2], optional
|
|
32
|
+
Order of finite difference gradients at domain edges, by default 2
|
|
33
|
+
|
|
34
|
+
Returns
|
|
35
|
+
-------
|
|
36
|
+
xr.DataArray
|
|
37
|
+
Directional component of Re[FT(u)* • FT((u•∇)u)]
|
|
38
|
+
"""
|
|
39
|
+
vel = (data.w, data.v, data.u)[axis.value]
|
|
40
|
+
adv_realspace = advection_xr(data, axis, method, edge_order=edge_order)
|
|
41
|
+
vel_hat = np.conj(xrft.fft(vel, dim=["x_caa", "y_aca", "z_aac"]))
|
|
42
|
+
adv_spec = xrft.fft(adv_realspace, dim=["x_caa", "y_aca", "z_aac"])
|
|
43
|
+
return np.conj(adv_spec * vel_hat) # type: ignore
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
def fourier_prep_xr(
|
|
47
|
+
data: xr.Dataset,
|
|
48
|
+
grad_method: GradMethod = GradMethod.numpy,
|
|
49
|
+
edge_order: Literal[1, 2] = 2,
|
|
50
|
+
) -> xr.Dataset:
|
|
51
|
+
"""Calculate integrand for spectral flux via fourier methods, with optional cupy
|
|
52
|
+
acceleration
|
|
53
|
+
|
|
54
|
+
Parameters
|
|
55
|
+
----------
|
|
56
|
+
data : xr.Dataset
|
|
57
|
+
Dataset with realspace velocity components (on the same grid)
|
|
58
|
+
grad_method : GradMethod, optional
|
|
59
|
+
Method for calculating gradients, by default GradMethod.numpy
|
|
60
|
+
edge_order : Literal[1, 2], optional
|
|
61
|
+
Order of finite difference gradients at domain edges, by default 2
|
|
62
|
+
|
|
63
|
+
Returns
|
|
64
|
+
-------
|
|
65
|
+
xr.Dataset
|
|
66
|
+
Dataset with integrand values, k magnitudes
|
|
67
|
+
"""
|
|
68
|
+
pi_int = (
|
|
69
|
+
pi_int_dir_xr(data, Axis.x, grad_method, edge_order)
|
|
70
|
+
+ pi_int_dir_xr(data, Axis.y, grad_method, edge_order)
|
|
71
|
+
+ pi_int_dir_xr(data, Axis.z, grad_method, edge_order)
|
|
72
|
+
).rename("pi_int")
|
|
73
|
+
k = (pi_int.freq_x_caa**2 + pi_int.freq_y_aca**2 + pi_int.freq_z_aac**2).rename("k")
|
|
74
|
+
out = xr.merge([pi_int, k]).assign_attrs({"L": data.x_caa[-1] - data.x_caa[0]})
|
|
75
|
+
if pi_int.cupy.is_cupy:
|
|
76
|
+
return out.as_cupy()
|
|
77
|
+
else:
|
|
78
|
+
return out
|
|
79
|
+
|
|
80
|
+
|
|
81
|
+
def fourier_int_xr(data: xr.Dataset, klim: float) -> xr.DataArray:
|
|
82
|
+
"""Calculate spectral flux of energy dissipation through a fourier transform
|
|
83
|
+
at a specific wavenumber
|
|
84
|
+
|
|
85
|
+
Parameters
|
|
86
|
+
----------
|
|
87
|
+
data : xr.Dataset
|
|
88
|
+
Dataset containing integrand, k magnitudes
|
|
89
|
+
klim : float
|
|
90
|
+
K value to integrate up to
|
|
91
|
+
|
|
92
|
+
Returns
|
|
93
|
+
-------
|
|
94
|
+
xr.DataArray
|
|
95
|
+
Spectral flux at klim
|
|
96
|
+
"""
|
|
97
|
+
masked = data["pi_int"].where(data["k"] <= (klim / (2 * np.pi)) ** 2, 0.0)
|
|
98
|
+
val = masked.integrate(["freq_x_caa", "freq_y_aca", "freq_z_aac"])
|
|
99
|
+
num = np.real(val.item()) / data.L**3
|
|
100
|
+
return xr.DataArray(num, {"time": data.time, "k": klim})
|
|
101
|
+
|
|
102
|
+
|
|
103
|
+
def fourier_int_xr_lst(data: xr.Dataset, k_lst: Iterable[float]) -> xr.DataArray:
|
|
104
|
+
"""Calculate spectral flux of energy dissipation through a fourier transform
|
|
105
|
+
over a list of wavenumbers
|
|
106
|
+
|
|
107
|
+
Parameters
|
|
108
|
+
----------
|
|
109
|
+
data : xr.Dataset
|
|
110
|
+
Dataset containing integrand, k magnitudes
|
|
111
|
+
k_lst : Iterable[float]
|
|
112
|
+
List of k values to integrate to
|
|
113
|
+
|
|
114
|
+
Returns
|
|
115
|
+
-------
|
|
116
|
+
xr.DataArray
|
|
117
|
+
Spectral flux as a function of k, at given k values
|
|
118
|
+
"""
|
|
119
|
+
return xr.concat([fourier_int_xr(data, k) for k in k_lst], "k")
|
|
120
|
+
|
|
121
|
+
|
|
122
|
+
def van_atta_prep(
|
|
123
|
+
data: xr.Dataset, cor_ax: Axis, vel_ax: Axis, periodic: bool = True
|
|
124
|
+
) -> xr.DataArray:
|
|
125
|
+
"""Triple product calculation of transfer function using Van Atta and Chen (1969)
|
|
126
|
+
velocities are assumed to be: u along x_caa, v along y_aca, w along z_aac
|
|
127
|
+
|
|
128
|
+
Parameters
|
|
129
|
+
----------
|
|
130
|
+
data : xr.Dataset
|
|
131
|
+
n-dimensional velocity data in all directions, must include the cor_ax
|
|
132
|
+
but need not include other dimensions
|
|
133
|
+
cor_ax : Axis
|
|
134
|
+
axis along which to take the correlations
|
|
135
|
+
vel_ax : Axis
|
|
136
|
+
second axis from which to use velocities (may be the same as cor_ax)
|
|
137
|
+
periodic : bool
|
|
138
|
+
whether to rotate (if periodic) or shift (if not) the velocities along cor_ax
|
|
139
|
+
|
|
140
|
+
Returns
|
|
141
|
+
-------
|
|
142
|
+
xr.DataArray
|
|
143
|
+
transfer function as a function of k
|
|
144
|
+
"""
|
|
145
|
+
cor_ax_name = axis_name(cor_ax)
|
|
146
|
+
cor_ax_xr = data[cor_ax_name]
|
|
147
|
+
cor_vel = data[("w", "v", "u")[cor_ax.value]]
|
|
148
|
+
shift_vel = data[("w", "v", "u")[vel_ax.value]]
|
|
149
|
+
n_shifts = len(cor_ax_xr) // 2 if periodic else len(cor_ax_xr)
|
|
150
|
+
shift_range = range(-n_shifts, n_shifts + 1)
|
|
151
|
+
dx = cor_ax_xr[1] - cor_ax_xr[0]
|
|
152
|
+
if periodic:
|
|
153
|
+
si1i = xr.concat(
|
|
154
|
+
[
|
|
155
|
+
(shift_vel * cor_vel * shift_vel.roll({cor_ax_name: i}))
|
|
156
|
+
.mean(cor_ax_name)
|
|
157
|
+
.expand_dims({"r": [i * dx]})
|
|
158
|
+
for i in shift_range
|
|
159
|
+
],
|
|
160
|
+
"r",
|
|
161
|
+
)
|
|
162
|
+
else:
|
|
163
|
+
si1i = xr.concat(
|
|
164
|
+
[
|
|
165
|
+
(shift_vel * cor_vel * shift_vel.shift({cor_ax_name: i}))
|
|
166
|
+
.mean(cor_ax_name, skipna=True)
|
|
167
|
+
.expand_dims({"r": [i * dx]})
|
|
168
|
+
for i in shift_range
|
|
169
|
+
],
|
|
170
|
+
"r",
|
|
171
|
+
)
|
|
172
|
+
si1i = 0.5 * (si1i - si1i.isel(r=slice(None, None, -1)).data)
|
|
173
|
+
l1 = xrft.fft(si1i, dim="r", real_dim="r").as_numpy()
|
|
174
|
+
l1 = l1.assign_coords(k=("freq_r", l1["freq_r"].data * 2 * np.pi))
|
|
175
|
+
l1 = l1 * 1j * l1["k"] / (2 * np.pi)
|
|
176
|
+
return 4 * l1 - 2 * l1["k"] * l1.differentiate("k")
|
|
177
|
+
|
|
178
|
+
|
|
179
|
+
def van_atta_int(data: xr.DataArray, mean_axes: Collection[str]) -> xr.DataArray:
|
|
180
|
+
out: xr.DataArray = cumulative_simpson(data, coord="k") # type: ignore
|
|
181
|
+
mean = out.isel(freq_r=slice(None, -1)).mean(mean_axes).real
|
|
182
|
+
norm = mean - mean.isel(freq_r=-1)
|
|
183
|
+
return norm.assign_coords(k=("freq_r", mean["freq_r"].data * 2 * np.pi))
|