def flux_coordinates_and_axes(draw,
domain=((0.0, 1.0), (0.0, 1.0), (0.0, 1.0)),
min_side=2,
max_side=12):
"""Generates :py:class:`indica.converters.FluxSurfaceCoordinates` objects,
along with axes.
Parameters
----------
min_side : integer
The minimum number of elements in an unaligned dimension for the
default coordinate arrays. (Not available for all coordinate systems.)
max_side : integer
The maximum number of elements in an unaligned dimension for the
default coordinate arrays. (Not available for all coordinate systems.)
"""
transform = draw(flux_coordinates(domain))
x1_vals, x2_vals, t_vals = map(
np.ravel,
draw(
basis_coordinates((0.0, 0.0, 0.0), (1.0, 2 * np.pi, 120.0),
min_side, max_side)),
)
return (
transform,
coord_array(x1_vals, transform.x1_name),
coord_array(x2_vals, transform.x2_name),
coord_array(t_vals, "t"),
)
def test_spline_transect_coords_x1():
spline = Spline(func(R_grid, 0, t_grid), "R", trivial, "natural")
indices = [1, 3, 5]
alpha = coord_array(indices, coords.x1_name)
alpha_z_offset = coord_array([0.2, 0.1], coords.x2_name)
result = spline(coords, alpha, alpha_z_offset, t_grid)
assert_allclose(result, func(R_positions[indices], 0.0, t_grid))
def test_broadcast_spline_2d():
x_interp = utilities.coord_array(np.linspace(0.2, 0.5, 3),
"alpha") * utilities.coord_array(
np.linspace(0.5, 1.0, 4), "beta")
result = utilities.broadcast_spline(spline, spline_dims, spline_coords,
x_interp)
assert_allclose(result, func_3d(x_interp, y, t))
assert result.dims == ("alpha", "beta", "y", "t")
def trivial_transforms_and_axes(draw,
domain=((0.0, 1.0), (0.0, 1.0), (0.0, 1.0)),
min_side=2,
max_side=12):
transform = draw(trivial_transforms(domain))
min_vals, max_vals = zip(*domain)
x1, x2, t = map(
np.ravel,
draw(basis_coordinates(min_vals, max_vals, min_side, max_side)))
return (
transform,
coord_array(x1, transform.x1_name),
coord_array(x2, transform.x2_name),
coord_array(t, "t"),
)
def test_spline_fit():
flux_coords = FluxSurfaceCoordinates("poloidal")
flux_coords.set_equilibrium(fake_equilib)
knot_locations = [0.0, 0.5, 0.8, 1.05]
knot_vals = (DataArray([1.0, 0.9, 0.6, 0.0],
coords=[("rho_poloidal", knot_locations)]) +
0.05 * t_grid)
knot_vals.loc[knot_locations[-1]] = 0.0
expected_spline = Spline(knot_vals, "rho_poloidal", flux_coords)
input_vals = expected_spline(coords, R_positions.coords["alpha"], 0.0,
t_grid).assign_coords(alpha_z_offset=0)
input_vals.attrs["datatype"] = ("temperature", "electrons")
fitter = SplineFit(knot_locations, sess=MagicMock())
result_locations = coord_array(np.linspace(0, 1.05, 10), "rho_poloidal")
result, spline_fit, binned_input = fitter(result_locations, t_grid,
input_vals)
assert_allclose(input_vals, binned_input)
assert_allclose(spline_fit, knot_vals.transpose(*spline_fit.dims))
assert_allclose(
result, expected_spline(flux_coords, result_locations, 0.0, t_grid))
assert result.attrs["datatype"] == input_vals.attrs["datatype"]
assert isinstance(result.attrs["splines"], Spline)
assert result.attrs["transform"] == flux_coords
assert "provenance" in result.attrs
assert "provenance" in spline_fit.attrs
assert "provenance" in binned_input.attrs
def magnetic_coordinates(draw,
domain=((0.0, 1.0), (0.0, 1.0), (0.0, 1.0)),
name=None):
z, name, machine_dims = draw(magnetic_coordinate_arguments(domain, name))
if domain:
Rmin = domain[0][0]
else:
Rmin = machine_dims[0][0]
if Rmin < 0:
Bcoeff = draw(floats(min(0.1, -0.01 / Rmin), min(10.0, -0.99 / Rmin)))
else:
Bcoeff = draw(floats(0.1, 10.0))
B_alpha = draw(floats(0.0, 1.0))
equilib = draw(
fake_equilibria(
draw(floats(*domain[0])),
draw(floats(*domain[1])),
coord_array(
draw(monotonic_series(*domain[2], draw(integers(2, 20))), ),
"t",
),
Btot_alpha=B_alpha,
Btot_b=Bcoeff,
))
result = MagneticCoordinates(z, name, machine_dims)
result.set_equilibrium(equilib)
return result
def los_coordinates_and_axes(
draw,
domain=None,
min_los=2,
max_los=10,
domain_as_dims=False,
toroidal_skew=None,
name=None,
):
"""Generates :py:class:`indica.converters.LinesOfSightTransform` objects
with lines of sight radiating from a point.
At present this point is on the edge of the Tokamak, for reasons of
convenience.
Parameters
----------
domain: Tuple[Tuple[float, float], Tuple[float, float]], optional
A tuple giving the range of R-z values for which the transform is
guaranteed to work: ``((Rmin, Rmax), (zmin, zmax)``. If absent will
draw values.
min_los: int
The minimum number of lines of sight
max_los: int
The maximum number of lines of sight
domain_as_dims: bool
If True, use the domain as the machine dimensions.
toroidal_skew: Optional[bool]
Whether to include any toroidal angle to the lines of sight. Defaults
to drawing a value to decide.
name: Optional[str]
The name of the camera these lines of sight are for.
Returns
-------
transform: TransectCoordinates
The coordinate transform object.
machine_dims: Tuple[Tuple[float, float], Tuple[float, float]]
A tuple giving the boundaries of the Tokamak in R-z space:
``((Rmin, Rmax), (zmin, zmax)``.
"""
transform = draw(
los_coordinates(
domain,
min_los,
max_los,
domain_as_dims,
toroidal_skew,
name,
))
x1 = coord_array(np.arange(len(transform.x_start)), transform.x1_name)
x2 = DataArray(0)
t = DataArray(0)
return transform, x1, x2, t
def test_coord_array(vals, name):
coords = utilities.coord_array(vals, name)
np.testing.assert_array_equal(coords.data, vals)
np.testing.assert_array_equal(coords.coords[name], vals)
assert len(coords.dims) == 1
datatype = coords.attrs["datatype"]
if name == "R":
assert datatype[0] == "major_rad"
elif name == "t":
assert datatype[0] == "time"
elif name == "rho_poloidal":
assert datatype[0] == "norm_flux_pol"
elif name == "rho_toroidal":
assert datatype[0] == "norm_flux_tor"
else:
assert datatype[0] == name
assert datatype[1] == "plasma"
def magnetic_coordinates_and_axes(
draw,
domain=((0.0, 1.0), (0.0, 1.0), (0.0, 1.0)),
name=None,
min_vals=2,
max_vals=5,
):
transform = draw(magnetic_coordinates(domain, name))
width = transform.right - transform.left
R_start = transform.right - draw(floats(0.001, 0.5 * width))
R_end = transform.left + draw(floats(0.001, 0.5 * width))
if R_start - R_end < 0.1 * width:
R_start = +0.1 * width
n = draw(integers(min_vals, max_vals))
R = DataArray(np.linspace(R_start, R_end, n))
x1 = coord_array(
transform.convert_from_Rz(R, transform.z_los, 0)[0].data,
transform.x1_name)
x2 = DataArray(0)
t = DataArray(0)
return transform, x1, x2, t
def test_broadcast_spline_new_coord():
x_interp = utilities.coord_array(np.linspace(0.1, 0.4, 4), "x_new")
result = utilities.broadcast_spline(spline, spline_dims, spline_coords,
x_interp)
assert_allclose(result, func_3d(x_interp, y, t))
assert result.dims == ("x_new", "y", "t")
elif name == "t":
assert datatype[0] == "time"
elif name == "rho_poloidal":
assert datatype[0] == "norm_flux_pol"
elif name == "rho_toroidal":
assert datatype[0] == "norm_flux_tor"
else:
assert datatype[0] == name
assert datatype[1] == "plasma"
def func_3d(x, y, t):
return 2 * x + 3 * y - 0.5 * t
x = utilities.coord_array(np.linspace(0.0, 1.0, 5), "x")
y = utilities.coord_array(np.linspace(0.0, 10.0, 4), "y")
t = utilities.coord_array(np.linspace(50.0, 55.0, 6), "t")
spline_dims = ("y", "t")
spline_coords = {"y": y, "t": t}
interp_data = func_3d(x, y, t)
spline = scipy.interpolate.CubicSpline(x, interp_data, 0, "natural")
def test_broadcast_spline_new_coord():
x_interp = utilities.coord_array(np.linspace(0.1, 0.4, 4), "x_new")
result = utilities.broadcast_spline(spline, spline_dims, spline_coords,
x_interp)
assert_allclose(result, func_3d(x_interp, y, t))
assert result.dims == ("x_new", "y", "t")
import getpass
import itertools
import matplotlib.pyplot as plt
import numpy as np
from indica.equilibrium import Equilibrium
from indica.operators import InvertRadiation
from indica.readers import PPFReader
from indica.utilities import coord_array
cameras = ["v"]
R = coord_array(np.linspace(1.83, 3.9, 25), "R")
z = coord_array(np.linspace(-1.75, 2.0, 25), "z")
t = coord_array(np.linspace(45, 50, 5), "t")
reader = PPFReader(90279, 45.0, 50.0)
if reader.requires_authentication:
user = input("JET username: "******"JET password: "******"jetppf", "efit", 0)
hrts = reader.get("jetppf", "hrts", 0)
sxr = reader.get("jetppf", "sxr", 0)
equilibrium = Equilibrium(equilib_dat) # , hrts["te"])
# *********************************************************************
for data in itertools.chain(hrts.values(), equilib_dat.values(), sxr.values()):
from indica.converters.time import strip_provenance
from indica.utilities import coord_array
from .test_abstract_transform import coordinate_transforms_and_axes
from ..data_strategies import data_arrays_from_coords
from ..strategies import polynomial_functions
from ..strategies import sane_floats
from ..strategies import separable_functions
from ..strategies import smooth_functions
start_times = floats(50.0, 120.0)
end_times = start_times.map(lambda t: 120.0 - (t - 50.0))
samples = integers(2, 100)
methods = sampled_from(
["nearest", "zero", "linear", "slinear", "quadratic", "cubic"])
t_axes = builds(np.linspace, just(50.0), just(120.0),
integers(4, 50)).map(lambda t: coord_array(t, "t"))
@composite
def useful_data_arrays(
draw,
rel_sigma=0.02,
abs_sigma=1e-3,
data_strategy=separable_functions(
smooth_functions(max_val=1e3),
smooth_functions(max_val=1e3),
smooth_functions(max_val=1e3),
),
):
times = draw(t_axes)
transform, x1, x2, _ = draw(coordinate_transforms_and_axes(min_side=2))
import getpass
import itertools
import matplotlib.pyplot as plt
import numpy as np
from indica.equilibrium import Equilibrium
from indica.operators import SplineFit
from indica.readers import PPFReader
from indica.utilities import coord_array
rho = coord_array(np.linspace(0, 1.04, 100), "rho_poloidal")
t = coord_array(np.linspace(45, 50, 11), "t")
reader = PPFReader(90279, 45.0, 50.0)
if reader.requires_authentication:
user = input("JET username: "******"JET password: "******"jetppf", "efit", 0)
hrts = reader.get("jetppf", "hrts", 0)
lidr = reader.get("jetppf", "lidr", 0)
equilibrium = Equilibrium(equilib_dat)
# *********************************************************************
for data in itertools.chain(hrts.values(), lidr.values()):
data.indica.equilibrium = equilibrium
# Test the spline class
fake_equilib = FakeEquilibrium(0.6, 0.0)
R_positions = DataArray(np.linspace(1.0, 0.1, 10),
coords=[
("alpha", np.arange(10))
]).assign_attrs(datatype=("major_rad", "plasma"))
z_positions = DataArray(np.linspace(-0.1, 0.2, 10),
coords=[("alpha", np.arange(10))
]).assign_attrs(datatype=("z", "plasma"))
coords = TransectCoordinates(R_positions, z_positions)
coords.set_equilibrium(fake_equilib)
trivial = TrivialTransform()
R_grid = coord_array(np.linspace(0.0, 1.2, 7), "R")
z_grid = coord_array(np.linspace(-1.0, 1.0, 5), "z")
t_grid = coord_array(np.linspace(0.0, 10.0, 11), "t")
def func(R, z, t):
return R - z + 0.5 * t
def test_spline_trivial_coords_R():
spline = Spline(func(R_grid, z_grid, t_grid), "R", trivial, "natural")
R = coord_array(np.linspace(0.2, 0.8, 4), "R")
result = spline(trivial, R, z_grid, t_grid)
assert_allclose(result, func(R, z_grid, t_grid))
def test_spline_trivial_coords_z():
spline = Spline(func(R_grid, z_grid, t_grid), "z", trivial, "natural")
z = coord_array(np.linspace(-0.2, 0.2, 5), "z")
result = spline(trivial, R_grid, z, t_grid)
assert_allclose(result, func(R_grid, z, t_grid).transpose(*result.dims))
def test_spline_trivial_coords_R():
spline = Spline(func(R_grid, z_grid, t_grid), "R", trivial, "natural")
R = coord_array(np.linspace(0.2, 0.8, 4), "R")
result = spline(trivial, R, z_grid, t_grid)
assert_allclose(result, func(R, z_grid, t_grid))
