"""Implementation of a 3D toroidally closed polygon class."""
import numpy as np
import h5netcdf
import xarray as xr
from .polygon_2d_m import Polygon2D
from torx.autogrid_decorators_m import autogrid_method
from torx.autodoc_decorators_m import autodoc_class
[docs]
@autodoc_class
class Polygon3D:
"""3D polygon toroidally closed."""
[docs]
@classmethod
def initialize_from_file(self, file_name, variable_name, **kwargs):
"""Initialize the 3D polygon from a NetCDF file and variable name."""
self.file_name = file_name
self.variable_name = variable_name
ds = h5netcdf.File(file_name)
polygon_vertices = ds.variables[variable_name][:].data
phi_array = ds.variables["phi_array"][:].data
ds.close()
return self(phi_array, polygon_vertices, **kwargs)
[docs]
def __init__(self, phi_array, polygon_vertices,
interpolation_method="linear", num_field_periods=None):
"""
Initialize the 3D polygon with an array of toroidal angles.
A phi values 2D polygons are defined in RZ. It is also
possible to initialize a "3D" polygon with only a single phi value,
essentially creating a 2D (axisymmetric) polygon.
Since the 3D polygon is defined discretely in phi, an interpolation
method in phi is required. Current options are "linear".
If num_field_periods is passed, this is used to verify that the phi
array is evenly spaced in a range [0, 2pi/num_field_periods).
"""
self.phi_array = np.atleast_1d(phi_array)
self.polygon_vertices = polygon_vertices
self.n_planes = polygon_vertices.shape[0]
self.n_vertices = polygon_vertices.shape[-1]
assert self.polygon_vertices.shape == (self.n_planes, 2,
self.n_vertices)
assert self.phi_array.size == self.n_planes
if self.n_planes > 1:
self.dphi = self.phi_array[1] - self.phi_array[0]
if num_field_periods is not None:
self.phi_max = 2.0 * np.pi / num_field_periods
assert np.isclose(self.phi_array[-1] + self.dphi, self.phi_max)
else:
self.phi_max = self.phi_array[-1] + self.dphi
assert np.allclose(self.phi_array,
np.linspace(0.0, self.phi_max, num=self.n_planes,
endpoint=False))
else:
self.dphi = 0.0
self.phi_max = 0.0
if interpolation_method == "linear":
self._interpolate_phi = self._interpolate_phi_linear
else:
raise Exception((f"Interpolation method {interpolation_method} is "
"not supported!"))
def _interpolate_phi_linear(self, phi):
"""
Return a polygon at the given phi.
Linearly interpolates the polygon array in phi.
"""
polygon_vertices = np.zeros((2, self.n_vertices))
if self.n_planes == 1:
# No interpolation necessary, the 3D polygon is axisymmetric
polygon_vertices[:] = self.polygon_vertices[0, :, :]
else:
phi_mod = phi % self.phi_max
plane_left = int(phi_mod // self.dphi)
plane_right = (plane_left + 1) % self.n_planes
interp_frac = (phi_mod - self.phi_array[plane_left]) / self.dphi
step = self.polygon_vertices[plane_right, :, :] \
- self.polygon_vertices[plane_left, :, :]
polygon_vertices[:] = self.polygon_vertices[plane_left, :, :] \
+ step * interp_frac
interpolated_polygon = Polygon2D(x_points=polygon_vertices[0, :],
y_points=polygon_vertices[1, :])
return interpolated_polygon
[docs]
@autogrid_method
def points_inside(self, r_norm, z_norm, phi, **kwargs):
"""
Test whether the points are inside the polygon array.
Todo
----
This currently only works with xarray-type coordinates, because
the autogrid decorator has not been fully ported to 3D!
"""
RR, ZZ, _ = xr.broadcast(r_norm, z_norm, phi)
inside = xr.zeros_like(RR, dtype=bool)
for k in range(phi.size):
poly_interp = self._interpolate_phi(phi.values[k])
inside[{"phi":k}] = poly_interp.points_inside(RR[{"phi":k}].values,
ZZ[{"phi":k}].values)
return inside
[docs]
def point_inside(self, r_norm, z_norm, phi, **kwargs):
"""Test whether the single point is inside the polygon array."""
poly_interp = self._interpolate_phi(phi)
inside = poly_interp.point_inside(r_norm, z_norm)
return inside
