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.
@@ -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,2 @@
1
+ # spec_flux
2
+ 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,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))