def equilibrium_dat():
machine_dims = ((1.83, 3.9), (-1.75, 2.0))
start_time, end_time = 75.0, 80.0
Btot_factor = None
result = {}
nspace = 8
ntime = 3
times = np.linspace(start_time - 0.5, end_time + 0.5, ntime)
tfuncs = smooth_funcs((start_time, end_time), 0.01)
r_centre = (machine_dims[0][0] + machine_dims[0][1]) / 2
z_centre = (machine_dims[1][0] + machine_dims[1][1]) / 2
raw_result = {}
attrs = {
"transform": TrivialTransform(),
"provenance": MagicMock(),
"partial_provenance": MagicMock(),
}
result["rmag"] = xr.DataArray(r_centre + tfuncs(times),
coords=[("t", times)],
name="rmag",
attrs=attrs)
result["rmag"].attrs["datatype"] = ("major_rad", "mag_axis")
result["zmag"] = xr.DataArray(z_centre + tfuncs(times),
coords=[("t", times)],
name="zmag",
attrs=attrs)
result["zmag"].attrs["datatype"] = ("z", "mag_axis")
fmin = 0.1
result["faxs"] = xr.DataArray(
fmin + np.abs(tfuncs(times)),
{
"t": times,
"R": result["rmag"],
"z": result["zmag"]
},
["t"],
name="faxs",
attrs=attrs,
)
result["faxs"].attrs["datatype"] = ("magnetic_flux", "mag_axis")
a_coeff = xr.DataArray(
np.vectorize(lambda x: 0.8 * x)(np.minimum(
np.abs(machine_dims[0][0] - result["rmag"]),
np.abs(machine_dims[0][1] - result["rmag"]),
), ),
coords=[("t", times)],
)
if Btot_factor is None:
b_coeff = xr.DataArray(
np.vectorize(lambda x: 0.8 * x)(np.minimum(
np.abs(machine_dims[1][0] - result["zmag"].data),
np.abs(machine_dims[1][1] - result["zmag"].data),
), ),
coords=[("t", times)],
)
n_exp = 0.5
fmax = 5.0
result["fbnd"] = xr.DataArray(
fmax - np.abs(tfuncs(times)),
coords=[("t", times)],
name="fbnd",
attrs=attrs,
)
else:
b_coeff = a_coeff
n_exp = 1
fdiff_max = Btot_factor * a_coeff
result["fbnd"] = xr.DataArray(
np.vectorize(lambda axs, diff: axs + 0.03 * diff)(
result["faxs"], fdiff_max.values),
coords=[("t", times)],
name="fbnd",
attrs=attrs,
)
result["fbnd"].attrs["datatype"] = ("magnetic_flux", "separtrix")
thetas = xr.DataArray(np.linspace(0.0, 2 * np.pi, nspace, endpoint=False),
dims=["arbitrary_index"])
result["rbnd"] = (
result["rmag"] + a_coeff * b_coeff /
np.sqrt(a_coeff**2 * np.tan(thetas)**2 + b_coeff**2)).assign_attrs(
**attrs)
result["rbnd"].name = "rbnd"
result["rbnd"].attrs["datatype"] = ("major_rad", "separatrix")
result["zbnd"] = (
result["zmag"] + a_coeff * b_coeff /
np.sqrt(a_coeff**2 + b_coeff**2 * np.tan(thetas)**-2)).assign_attrs(
**attrs)
result["zbnd"].name = "zbnd"
result["zbnd"].attrs["datatype"] = ("z", "separatrix")
r = np.linspace(machine_dims[0][0], machine_dims[0][1], nspace)
z = np.linspace(machine_dims[1][0], machine_dims[1][1], nspace)
rgrid = xr.DataArray(r, coords=[("R", r)])
zgrid = xr.DataArray(z, coords=[("z", z)])
psin = ((-result["zmag"] + zgrid)**2 / b_coeff**2 +
(-result["rmag"] + rgrid)**2 / a_coeff**2)**(0.5 / n_exp)
psi = psin * (result["fbnd"] - result["faxs"]) + result["faxs"]
psi.name = "psi"
psi.attrs["transform"] = attrs["transform"]
psi.attrs["provenance"] = MagicMock()
psi.attrs["partial_provenance"] = MagicMock()
psi.attrs["datatype"] = ("magnetic_flux", "plasma")
result["psi"] = psi
psin_coords = np.linspace(0.0, 1.0, nspace)
rho = np.sqrt(psin_coords)
psin_data = xr.DataArray(psin_coords, coords=[("rho_poloidal", rho)])
attrs["transform"] = FluxSurfaceCoordinates("poloidal", )
def monotonic_series(start,
stop,
num=50,
endpoint=True,
retstep=False,
dtype=None):
return np.linspace(start, stop, num, endpoint, retstep, dtype)
ftor_min = 0.1
ftor_max = 5.0
result["ftor"] = xr.DataArray(
np.outer(1 + tfuncs(times), monotonic_series(ftor_min, ftor_max,
nspace)),
coords=[("t", times), ("rho_poloidal", rho)],
name="ftor",
attrs=attrs,
)
result["ftor"].attrs["datatype"] = ("toroidal_flux", "plasma")
if Btot_factor is None:
f_min = 0.1
f_max = 5.0
time_vals = tfuncs(times)
space_vals = monotonic_series(f_min, f_max, nspace)
f_raw = np.outer(abs(1 + time_vals), space_vals)
else:
f_raw = np.outer(
np.sqrt(Btot_factor**2 -
(raw_result["fbnd"] - raw_result["faxs"])**2 / a_coeff**2),
np.ones_like(rho),
)
f_raw[:, 0] = Btot_factor
result["f"] = xr.DataArray(f_raw,
coords=[("t", times), ("rho_poloidal", rho)],
name="f",
attrs=attrs)
result["f"].attrs["datatype"] = ("f_value", "plasma")
result["rmjo"] = (result["rmag"] +
a_coeff * psin_data**n_exp).assign_attrs(**attrs)
result["rmjo"].name = "rmjo"
result["rmjo"].attrs["datatype"] = ("major_rad", "lfs")
result["rmjo"].coords["z"] = result["zmag"]
result["rmji"] = (result["rmag"] -
a_coeff * psin_data**n_exp).assign_attrs(**attrs)
result["rmji"].name = "rmji"
result["rmji"].attrs["datatype"] = ("major_rad", "hfs")
result["rmji"].coords["z"] = result["zmag"]
result["vjac"] = (4 * n_exp * np.pi**2 * result["rmag"] * a_coeff *
b_coeff *
psin_data**(2 * n_exp - 1)).assign_attrs(**attrs)
result["vjac"].name = "vjac"
result["vjac"].attrs["datatype"] = ("volume_jacobian", "plasma")
return result
def equilibrium_data(
draw,
machine_dims=((1.83, 3.9), (-1.75, 2.0)),
min_spatial_points=3,
max_spatial_points=12,
min_time_points=2,
max_time_points=15,
start_time=75.0,
end_time=80.0,
Btot_factor=None,
):
"""Returns a dictionary containing the data necessary to construct an
:py:class:`indica.equilibrium.Equilibrium` object.
Parameters
----------
machine_dims
The size of the reactor, ((Rmin, Rmax), (zmin, zmax))
min_spatial_points
The minimum number of points to use for spatial coordinate axes
max_spatial_points
The maximum number of points to use for spatial coordinate axes
min_time_points
The minimum number of points to use for the time axis
max_time_points
The maximum number of points to use for the time axis
Btot_factor
If present, the equilibrium will have total magnetic field strength
Btot_factor/R.
"""
result = {}
nspace = draw(integers(min_spatial_points, max_spatial_points))
ntime = draw(integers(min_time_points, max_time_points))
times = np.linspace(start_time - 0.5, end_time + 0.5, ntime)
tfuncs = smooth_functions((start_time, end_time), 0.01)
r_centre = (machine_dims[0][0] + machine_dims[0][1]) / 2
z_centre = (machine_dims[1][0] + machine_dims[1][1]) / 2
raw_result = {}
attrs = {
"transform": TrivialTransform(),
"provenance": MagicMock(),
"partial_provenance": MagicMock(),
}
result["rmag"] = DataArray(r_centre + draw(tfuncs)(times),
coords=[("t", times)],
name="rmag",
attrs=attrs)
result["rmag"].attrs["datatype"] = ("major_rad", "mag_axis")
result["zmag"] = DataArray(z_centre + draw(tfuncs)(times),
coords=[("t", times)],
name="zmag",
attrs=attrs)
result["zmag"].attrs["datatype"] = ("z", "mag_axis")
fmin = draw(floats(0.0, 1.0))
result["faxs"] = DataArray(
fmin + np.abs(draw(tfuncs)(times)),
{
"t": times,
"R": result["rmag"],
"z": result["zmag"]
},
["t"],
name="faxs",
attrs=attrs,
)
result["faxs"].attrs["datatype"] = ("magnetic_flux", "mag_axis")
a_coeff = DataArray(
np.vectorize(lambda x: draw(floats(0.75 * x, x)))(np.minimum(
np.abs(machine_dims[0][0] - result["rmag"]),
np.abs(machine_dims[0][1] - result["rmag"]),
)),
coords=[("t", times)],
)
if Btot_factor is None:
b_coeff = DataArray(
np.vectorize(lambda x: draw(floats(0.75 * x, x)))(np.minimum(
np.abs(machine_dims[1][0] - result["zmag"].data),
np.abs(machine_dims[1][1] - result["zmag"].data),
), ),
coords=[("t", times)],
)
n_exp = 0.5 # 1 # draw(floats(-0.5, 1.0))
fmax = draw(floats(max(1.0, 2 * fmin), 10.0))
result["fbnd"] = DataArray(
fmax - np.abs(draw(tfuncs)(times)),
coords=[("t", times)],
name="fbnd",
attrs=attrs,
)
else:
b_coeff = a_coeff
n_exp = 1
fdiff_max = Btot_factor * a_coeff
result["fbnd"] = DataArray(
np.vectorize(
lambda axs, diff: axs + draw(floats(0.001 * diff, diff)))(
result["faxs"], fdiff_max.values),
coords=[("t", times)],
name="fbnd",
attrs=attrs,
)
result["fbnd"].attrs["datatype"] = ("magnetic_flux", "separatrix")
thetas = DataArray(np.linspace(0.0, 2 * np.pi, nspace, endpoint=False),
dims=["arbitrary_index"])
result["rbnd"] = (
result["rmag"] + a_coeff * b_coeff /
np.sqrt(a_coeff**2 * np.tan(thetas)**2 + b_coeff**2)).assign_attrs(
**attrs)
result["rbnd"].name = "rbnd"
result["rbnd"].attrs["datatype"] = ("major_rad", "separatrix")
result["zbnd"] = (
result["zmag"] + a_coeff * b_coeff /
np.sqrt(a_coeff**2 + b_coeff**2 * np.tan(thetas)**-2)).assign_attrs(
**attrs)
result["zbnd"].name = "zbnd"
result["zbnd"].attrs["datatype"] = ("z", "separatrix")
r = np.linspace(machine_dims[0][0], machine_dims[0][1], nspace)
z = np.linspace(machine_dims[1][0], machine_dims[1][1], nspace)
rgrid = DataArray(r, coords=[("R", r)])
zgrid = DataArray(z, coords=[("z", z)])
psin = ((-result["zmag"] + zgrid)**2 / b_coeff**2 +
(-result["rmag"] + rgrid)**2 / a_coeff**2)**(0.5 / n_exp)
psi = psin * (result["fbnd"] - result["faxs"]) + result["faxs"]
psi.name = "psi"
psi.attrs["transform"] = attrs["transform"]
psi.attrs["provenance"] = MagicMock()
psi.attrs["partial_provenance"] = MagicMock()
psi.attrs["datatype"] = ("magnetic_flux", "plasma")
result["psi"] = psi
psin_coords = np.linspace(0.0, 1.0, nspace)
rho = np.sqrt(psin_coords)
psin_data = DataArray(psin_coords, coords=[("rho_poloidal", rho)])
attrs["transform"] = FluxSurfaceCoordinates("poloidal", )
ftor_min = draw(floats(0.0, 1.0))
ftor_max = draw(floats(max(1.0, 2 * fmin), 10.0))
result["ftor"] = DataArray(
np.outer(1 + draw(tfuncs)(times),
draw(monotonic_series(ftor_min, ftor_max, nspace))),
coords=[("t", times), ("rho_poloidal", rho)],
name="ftor",
attrs=attrs,
)
result["ftor"].attrs["datatype"] = ("toroidal_flux", "plasma")
if Btot_factor is None:
f_min = draw(floats(0.0, 1.0))
f_max = draw(floats(max(1.0, 2 * fmin), 10.0))
time_vals = draw(tfuncs)(times)
space_vals = draw(monotonic_series(f_min, f_max, nspace))
f_raw = np.outer(abs(1 + time_vals), space_vals)
else:
f_raw = np.outer(
np.sqrt(Btot_factor**2 -
(raw_result["fbnd"] - raw_result["faxs"])**2 / a_coeff**2),
np.ones_like(rho),
)
f_raw[:, 0] = Btot_factor
result["f"] = DataArray(f_raw,
coords=[("t", times), ("rho_poloidal", rho)],
name="f",
attrs=attrs)
result["f"].attrs["datatype"] = ("f_value", "plasma")
result["rmjo"] = (result["rmag"] +
a_coeff * psin_data**n_exp).assign_attrs(**attrs)
result["rmjo"].name = "rmjo"
result["rmjo"].attrs["datatype"] = ("major_rad", "lfs")
result["rmjo"].coords["z"] = result["zmag"]
result["rmji"] = (result["rmag"] -
a_coeff * psin_data**n_exp).assign_attrs(**attrs)
result["rmji"].name = "rmji"
result["rmji"].attrs["datatype"] = ("major_rad", "hfs")
result["rmji"].coords["z"] = result["zmag"]
result["vjac"] = (4 * n_exp * np.pi**2 * result["rmag"] * a_coeff *
b_coeff *
psin_data**(2 * n_exp - 1)).assign_attrs(**attrs)
result["vjac"].name = "vjac"
result["vjac"].attrs["datatype"] = ("volume_jacobian", "plasma")
return result
def data_arrays_from_coords(
draw,
data_type=(None, None),
coordinates=TrivialTransform(),
axes=(0.0, 0.0, 0.0),
data=separable_functions(
smooth_functions(max_val=1e3),
smooth_functions(max_val=1e3),
smooth_functions(max_val=1e3),
),
rel_sigma=0.02,
abs_sigma=1e-3,
uncertainty=True,
max_dropped=0.1,
require_dropped=False,
):
"""Returns a DataArray which uses the given coordinate transform.
Parameters
----------
data_type : Tuple[str, str]
The data type of the data_array to be generated. If either element of
the tuple is ``None`` then that element will be drawn from a strategy.
coordinates
A coordinate transform to use for this data.
axes
If item is not None, use those coordinates rather than the defaults
from the coordinate transform. Should be ordered ``[x1, x2, t]``.
data
A strategy to generate functions which calculate the contents of the
DataArray from the coordinates. Note that all coordinates will be
normalised before being passed to this function.
rel_sigma
Standard deviation of relative noise applied to the data
abs_sigma
Standard deviation of absolute noise applied to the data
uncertainty
If ``True``, generate uncertainty metadata using ``rel_sigma`` and
``abs_sigma`` (if they are non-zero).
max_dropped
The maximum number of channels to drop, as a fraction of the total
number of channels.
require_dropped
If True, ensure at least one channel is dropped.
"""
general_type = (data_type[0] if data_type[0] else draw(
general_datatypes(data_type[1])))
specific_type = (data_type[1] if data_type[1] else draw(
specific_datatypes(general_type)))
x1, x2, t = axes
func = (draw(noisy_functions(draw(data), rel_sigma, abs_sigma))
if rel_sigma or abs_sigma else draw(data))
coords = {}
dims = []
for n, c in [("t", t), (coordinates.x1_name, x1),
(coordinates.x2_name, x2)]:
if isinstance(c, np.ndarray) and c.ndim > 0:
coords[n] = c.flatten()
dims.append(n)
elif isinstance(c, DataArray) and c.ndim > 0:
coords[n] = c.values
dims.append(n)
elif not isinstance(coordinates, LinesOfSightTransform):
coords[n] = c
shape = tuple(len(coords[dim]) for dim in dims)
if isinstance(t, (np.ndarray, DataArray)) and t.ndim > 0:
min_val = np.min(t)
width = np.abs(np.max(t) - min_val)
t_scaled = (t - min_val) / (width if width else 1.0)
else:
t_scaled = 0.0
if isinstance(x1, (np.ndarray, DataArray)) and x1.ndim > 0:
min_val = np.min(x1)
width = np.abs(np.max(x1) - min_val)
x1_scaled = (x1 - min_val) / (width if width else 1.0)
else:
x1_scaled = 0.0
if isinstance(x2, (np.ndarray, DataArray)) and x2.ndim > 0:
min_val = np.min(x2)
width = np.abs(np.max(x2) - min_val)
x2_scaled = (x2 - min_val) / (width if width else 1.0)
else:
x2_scaled = 0.0
tmp = func(t_scaled, x1_scaled, x2_scaled)
result = DataArray(np.reshape(tmp, shape), coords, dims)
if isinstance(x1, np.ndarray):
flat_x1 = x1.flatten()
else:
flat_x1 = x1
dropped = ([
flat_x1[i] for i in draw(dropped_channels(len(x1), max_dropped))
] if isinstance(x1, np.ndarray) else [])
if require_dropped and len(dropped) == 0:
dropped = [flat_x1[0]]
if uncertainty and (rel_sigma or abs_sigma):
error = rel_sigma * result + abs_sigma
result.attrs["error"] = error
if dropped and flat_x1[0] != flat_x1[-1]:
to_keep = np.logical_not(
DataArray(flat_x1, coords=[("x1", flat_x1)]).isin(dropped))
dropped_result = result.sel(x1=dropped)
result = result.where(to_keep)
if uncertainty and (rel_sigma or abs_sigma):
dropped_result.attrs["error"] = result.attrs["error"].sel(
x1=dropped)
result.attrs["error"] = result.attrs["error"].where(to_keep)
result.attrs["dropped"] = dropped_result
result.attrs["datatype"] = (general_type, specific_type)
result.attrs["provenance"] = MagicMock()
result.attrs["partial_provenance"] = MagicMock()
result.attrs["transform"] = coordinates
if hasattr(coordinates, "equilibrium"):
result.indica.equilibrium = coordinates.equilibrium
else:
result.indica.equilibrium = MagicMock()
return result
# 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 data_arrays_from_coord(
data_type=(None, None),
coordinates=TrivialTransform(),
axes=(0.0, 0.0, 0.0),
data=separable_funcs(
smooth_funcs(max_val=1e3),
smooth_funcs(max_val=1e3),
smooth_funcs(max_val=1e3),
),
rel_sigma=0.02,
abs_sigma=1e-3,
uncertainty=True,
):
general_type = data_type[0]
specific_type = data_type[1]
x1, x2, t = axes
coords = {}
dims = []
for n, c in [("t", t), (coordinates.x1_name, x1),
(coordinates.x2_name, x2)]:
if isinstance(c, np.ndarray) and c.ndim > 0:
coords[n] = c.flatten()
dims.append(n)
elif isinstance(c, xr.DataArray) and c.ndim > 0:
coords[n] = c.values
dims.append(n)
elif not isinstance(coordinates, LinesOfSightTransform):
coords[n] = c if c is not None else 0.0
shape = tuple(len(coords[dim]) for dim in dims)
if isinstance(t, (np.ndarray, xr.DataArray)) and t.ndim > 0:
min_val = np.min(t)
width = np.abs(np.max(t) - min_val)
t_scaled = (t - min_val) / (width if width else 1.0)
else:
t_scaled = np.array([0.0])
if isinstance(x1, (np.ndarray, xr.DataArray)) and x1.ndim > 0:
min_val = np.min(x1)
width = np.abs(np.max(x1) - min_val)
x1_scaled = (x1 - min_val) / (width if width else 1.0)
else:
x1_scaled = np.array([0.0])
if isinstance(x2, (np.ndarray, xr.DataArray)) and x2.ndim > 0:
min_val = np.min(x2)
width = np.abs(np.max(x2) - min_val)
x2_scaled = (x2 - min_val) / (width if width else 1.0)
else:
x2_scaled = np.array([0.0])
tmp = data(x1_scaled, x2_scaled, t_scaled)
result = xr.DataArray(np.reshape(tmp, shape), coords, dims)
if uncertainty and (rel_sigma and abs_sigma):
error = rel_sigma * result + abs_sigma
result.attrs["error"] = error
result.attrs["datatype"] = (general_type, specific_type)
result.attrs["provenance"] = MagicMock()
result.attrs["partial_provenance"] = MagicMock()
result.attrs["transform"] = coordinates
if hasattr(coordinates, "equilibrium"):
result.indica.equilibrium = coordinates.equilibrium
else:
result.indica.equilibrium = MagicMock()
return result
def trivial_transforms(draw, domain=((0.0, 1.0), (0.0, 1.0), (0.0, 1.0))):
result = TrivialTransform()
result.set_equilibrium(MagicMock())
return result