def __edge(self, field, sx, sy, diff1, diff2, rec):
import numpy as np
from pencil.calc.streamlines import Stream
from pencil.math.interpolation import vec_int
phi_min = np.pi / 8.0
dtot = np.arctan2(
diff1[0] * diff2[1] - diff2[0] * diff1[1],
diff1[0] * diff2[0] + diff1[1] * diff2[1],
)
if (abs(dtot) > phi_min) and (rec < 4):
xm = 0.5 * (sx[0] + sx[1])
ym = 0.5 * (sy[0] + sy[1])
# Trace the intermediate field line.
# time = np.linspace(0, self.params.Lz/np.max(abs(field[2])), 100)
field_strength_z0 = vec_int(
np.array([xm, ym, self.params.Oz]),
field,
[self.params.dx, self.params.dy, self.params.dz],
[self.params.Ox, self.params.Oy, self.params.Oz],
[self.params.nx, self.params.ny, self.params.nz],
interpolation=self.params.interpolation,
)
field_strength_z0 = np.sqrt(np.sum(field_strength_z0 ** 2))
time = np.linspace(0, 4 * self.params.Lz / field_strength_z0, 500)
stream = Stream(
field, self.params, xx=np.array([xm, ym, self.params.Oz]), time=time
)
stream_x0 = stream.tracers[0, 0]
stream_y0 = stream.tracers[0, 1]
stream_x1 = stream.tracers[-1, 0]
stream_y1 = stream.tracers[-1, 1]
diffm = np.array([stream_x1 - stream_x0, stream_y1 - stream_y0])
if sum(diffm ** 2) != 0:
diffm = diffm / np.sqrt(sum(diffm ** 2))
dtot = self.__edge(
field, [sx[0], xm], [sy[0], ym], diff1, diffm, rec + 1
) + self.__edge(field, [xm, sx[1]], [ym, sy[1]], diffm, diff2, rec + 1)
return dtot
def __grad_field(self, xyz, var, field, dd):
"""
Compute the gradient if the field at xyz.
"""
import numpy as np
from pencil.math.interpolation import vec_int
gf = np.zeros((3, 3))
gf[0, :] = (vec_int(xyz + np.array([dd, 0, 0]), var, field) -
vec_int(xyz - np.array([dd, 0, 0]), var, field)) / (2 * dd)
gf[1, :] = (vec_int(xyz + np.array([0, dd, 0]), var, field) -
vec_int(xyz - np.array([0, dd, 0]), var, field)) / (2 * dd)
gf[2, :] = (vec_int(xyz + np.array([0, 0, dd]), var, field) -
vec_int(xyz - np.array([0, 0, dd]), var, field)) / (2 * dd)
return np.matrix(gf)
def __null_point(self, point, var, field):
import numpy as np
from pencil.calc.streamlines import Stream
from pencil.math.interpolation import vec_int
dl = np.min([var.dx, var.dy]) / 30.0
it = 0
# Tracers used to find the fixed point.
tracers_null = np.zeros((5, 4))
while True:
# Trace field lines at original point and for Jacobian.
# (second order seems to be enough)
xx = np.zeros((5, 3))
xx[0, :] = np.array([point[0], point[1], self.params.Oz])
xx[1, :] = np.array([point[0] - dl, point[1], self.params.Oz])
xx[2, :] = np.array([point[0] + dl, point[1], self.params.Oz])
xx[3, :] = np.array([point[0], point[1] - dl, self.params.Oz])
xx[4, :] = np.array([point[0], point[1] + dl, self.params.Oz])
for it1 in range(5):
# time = time = np.linspace(0, self.params.Lz/np.max(abs(field[2])), 10)
field_strength_z0 = vec_int(
xx[it1, :],
field,
[var.dx, var.dy, var.dz],
[var.x[0], var.y[0], var.z[0]],
[len(var.x), len(var.y), len(var.z)],
interpolation=self.params.interpolation,
)
field_strength_z0 = np.sqrt(np.sum(field_strength_z0 ** 2))
time = np.linspace(0, 4 * self.params.Lz / field_strength_z0, 500)
stream = Stream(field, self.params, xx=xx[it1, :], time=time)
tracers_null[it1, :2] = xx[it1, :2]
tracers_null[it1, 2:] = stream.tracers[-1, 0:2]
# Check function convergence.
ff = np.zeros(2)
ff[0] = tracers_null[0, 2] - tracers_null[0, 0]
ff[1] = tracers_null[0, 3] - tracers_null[0, 1]
if sum(abs(ff)) <= 1e-3 * np.min([self.params.dx, self.params.dy]):
fixed_point = np.array([point[0], point[1]])
break
# Compute the Jacobian.
fjac = np.zeros((2, 2))
fjac[0, 0] = (
(
(tracers_null[2, 2] - tracers_null[2, 0])
- (tracers_null[1, 2] - tracers_null[1, 0])
)
/ 2.0
/ dl
)
fjac[0, 1] = (
(
(tracers_null[4, 2] - tracers_null[4, 0])
- (tracers_null[3, 2] - tracers_null[3, 0])
)
/ 2.0
/ dl
)
fjac[1, 0] = (
(
(tracers_null[2, 3] - tracers_null[2, 1])
- (tracers_null[1, 3] - tracers_null[1, 1])
)
/ 2.0
/ dl
)
fjac[1, 1] = (
(
(tracers_null[4, 3] - tracers_null[4, 1])
- (tracers_null[3, 3] - tracers_null[3, 1])
)
/ 2.0
/ dl
)
# Invert the Jacobian.
fjin = np.zeros((2, 2))
det = fjac[0, 0] * fjac[1, 1] - fjac[0, 1] * fjac[1, 0]
if abs(det) < dl:
fixed_point = point
break
fjin[0, 0] = fjac[1, 1]
fjin[1, 1] = fjac[0, 0]
fjin[0, 1] = -fjac[0, 1]
fjin[1, 0] = -fjac[1, 0]
fjin = fjin / det
dpoint = np.zeros(2)
dpoint[0] = -fjin[0, 0] * ff[0] - fjin[0, 1] * ff[1]
dpoint[1] = -fjin[1, 0] * ff[0] - fjin[1, 1] * ff[1]
point += dpoint
# Check root convergence.
if sum(abs(dpoint)) < 1e-3 * np.min([self.params.dx, self.params.dy]):
fixed_point = point
print("Root finding converged.")
break
if it > 20:
fixed_point = point
print("Root finding did not converge.")
break
it += 1
return fixed_point
def find_fixed(
self,
datadir="data",
var_file="VAR0",
trace_field="bb",
ti=-1,
tf=-1,
tracer_file_name=None,
):
"""
Find the fixed points to a snapshot or existing tracer file.
call signature::
find_fixed(datadir='data', var_file='VAR0', trace_field='bb',
ti=-1, tf=-1, tracer_file_name=None):
Keyword arguments:
*datadir*:
Data directory.
*var_file*:
Varfile to be read.
*trace_field*:
Vector field used for the streamline tracing.
*ti*:
Initial VAR file index for tracer time sequences. Overrides 'var_file'.
*tf*:
Final VAR file index for tracer time sequences. Overrides 'var_file'.
*tracer_file_name*
Name of the tracer file to be read.
If 'None' compute the tracers.
"""
import numpy as np
import multiprocessing as mp
from pencil import read
from pencil import math
from pencil.diag.tracers import Tracers
from pencil.calc.streamlines import Stream
from pencil.math.interpolation import vec_int
if self.params.int_q == "curly_A":
self.curly_A = []
if self.params.int_q == "ee":
self.ee = []
# Multi core setup.
if not (np.isscalar(self.params.n_proc)) or (self.params.n_proc % 1 != 0):
print("Error: invalid processor number")
return -1
queue = mp.Queue()
# Make sure to read the var files with the correct magic.
magic = []
if trace_field == "bb":
magic.append("bb")
if trace_field == "jj":
magic.append("jj")
if trace_field == "vort":
magic.append("vort")
if self.params.int_q == "ee":
magic.append("bb")
magic.append("jj")
dim = read.dim(datadir=datadir)
# Check if user wants a tracer time series.
if (ti % 1 == 0) and (tf % 1 == 0) and (ti >= 0) and (tf >= ti):
series = True
var_file = "VAR{0}".format(ti)
n_times = tf - ti + 1
else:
series = False
n_times = 1
self.t = np.zeros(n_times)
# Read the initial field.
var = read.var(
var_file=var_file, datadir=datadir, magic=magic, quiet=True, trimall=True
)
self.t[0] = var.t
grid = read.grid(datadir=datadir, quiet=True, trim=True)
field = getattr(var, trace_field)
param2 = read.param(datadir=datadir, quiet=True)
if self.params.int_q == "ee":
ee = var.jj * param2.eta - math.cross(var.uu, var.bb)
self.params.datadir = datadir
self.params.var_file = var_file
self.params.trace_field = trace_field
# Get the simulation parameters.
self.params.dx = var.dx
self.params.dy = var.dy
self.params.dz = var.dz
self.params.Ox = var.x[0]
self.params.Oy = var.y[0]
self.params.Oz = var.z[0]
self.params.Lx = grid.Lx
self.params.Ly = grid.Ly
self.params.Lz = grid.Lz
self.params.nx = dim.nx
self.params.ny = dim.ny
self.params.nz = dim.nz
tracers = Tracers()
tracers.params = self.params
# Create the mapping for all times.
if not tracer_file_name:
tracers.find_tracers(
var_file=var_file,
datadir=datadir,
trace_field=trace_field,
ti=ti,
tf=tf,
)
else:
tracers.read(datadir=datadir, file_name=tracer_file_name)
self.tracers = tracers
# Set some default values.
self.t = np.zeros((tf - ti + 1) * series + (1 - series))
self.fixed_index = np.zeros((tf - ti + 1) * series + (1 - series))
self.poincare = np.zeros(
[
int(self.params.trace_sub * dim.nx),
int(self.params.trace_sub * dim.ny),
n_times,
]
)
ix0 = range(0, int(self.params.nx * self.params.trace_sub) - 1)
iy0 = range(0, int(self.params.ny * self.params.trace_sub) - 1)
self.fixed_points = []
self.fixed_sign = []
self.fixed_tracers = []
# Start the parallelized fixed point finding.
for tidx in range(n_times):
if tidx > 0:
var = read.var(
var_file="VAR{0}".format(tidx + ti),
datadir=datadir,
magic=magic,
quiet=True,
trimall=True,
)
field = getattr(var, trace_field)
self.t[tidx] = var.t
proc = []
sub_data = []
fixed = []
fixed_sign = []
fixed_tracers = []
for i_proc in range(self.params.n_proc):
proc.append(
mp.Process(
target=self.__sub_fixed,
args=(queue, ix0, iy0, field, self.tracers, tidx, var, i_proc),
)
)
for i_proc in range(self.params.n_proc):
proc[i_proc].start()
for i_proc in range(self.params.n_proc):
sub_data.append(queue.get())
for i_proc in range(self.params.n_proc):
proc[i_proc].join()
for i_proc in range(self.params.n_proc):
# Extract the data from the single cores. Mind the order.
sub_proc = sub_data[i_proc][0]
fixed.extend(sub_data[i_proc][1])
fixed_tracers.extend(sub_data[i_proc][2])
fixed_sign.extend(sub_data[i_proc][3])
self.fixed_index[tidx] += sub_data[i_proc][4]
self.poincare[sub_proc :: self.params.n_proc, :, tidx] = sub_data[
i_proc
][5]
for i_proc in range(self.params.n_proc):
proc[i_proc].terminate()
# Discard fixed points which lie too close to each other.
fixed, fixed_tracers, fixed_sign = self.__discard_close_fixed_points(
np.array(fixed), np.array(fixed_sign), np.array(fixed_tracers), var
)
if self.fixed_points is None:
self.fixed_points = []
self.fixed_sign = []
self.fixed_tracers = []
self.fixed_points.append(np.array(fixed))
self.fixed_sign.append(np.array(fixed_sign))
self.fixed_tracers.append(fixed_tracers)
# Compute the traced quantities along the fixed point streamlines.
if (self.params.int_q == "curly_A") or (self.params.int_q == "ee"):
for t_idx in range(0, n_times):
if self.params.int_q == "curly_A":
self.curly_A.append([])
if self.params.int_q == "ee":
self.ee.append([])
for fixed in self.fixed_points[t_idx]:
# Trace the stream line.
xx = np.array([fixed[0], fixed[1], self.params.Oz])
# time = np.linspace(0, self.params.Lz/np.max(abs(field[2])), 10)
field_strength_z0 = vec_int(
xx,
field,
[var.dx, var.dy, var.dz],
[var.x[0], var.y[0], var.z[0]],
[len(var.x), len(var.y), len(var.z)],
interpolation=self.params.interpolation,
)
field_strength_z0 = np.sqrt(np.sum(field_strength_z0 ** 2))
time = np.linspace(0, 4 * self.params.Lz / field_strength_z0, 500)
stream = Stream(field, self.params, xx=xx, time=time)
# Do the field line integration.
if self.params.int_q == "curly_A":
curly_A = 0
for l in range(stream.iterations - 1):
aaInt = vec_int(
(stream.tracers[l + 1] + stream.tracers[l]) / 2,
var.aa,
[var.dx, var.dy, var.dz],
[var.x[0], var.y[0], var.z[0]],
[len(var.x), len(var.y), len(var.z)],
interpolation=self.params.interpolation,
)
curly_A += np.dot(
aaInt, (stream.tracers[l + 1] - stream.tracers[l])
)
self.curly_A[-1].append(curly_A)
if self.params.int_q == "ee":
ee_p = 0
for l in range(stream.iterations - 1):
eeInt = vec_int(
(stream.tracers[l + 1] + stream.tracers[l]) / 2,
ee,
[var.dx, var.dy, var.dz],
[var.x[0], var.y[0], var.z[0]],
[len(var.x), len(var.y), len(var.z)],
interpolation=self.params.interpolation,
)
ee_p += np.dot(
eeInt, (stream.tracers[l + 1] - stream.tracers[l])
)
self.ee[-1].append(ee_p)
if self.params.int_q == "curly_A":
self.curly_A[-1] = np.array(self.curly_A[-1])
if self.params.int_q == "ee":
self.ee[-1] = np.array(self.ee[-1])
return 0
def __sub_fixed(self, queue, ix0, iy0, field, tracers, tidx, var, i_proc):
import numpy as np
from pencil.calc.streamlines import Stream
from pencil.math.interpolation import vec_int
diff = np.zeros((4, 2))
fixed = []
fixed_sign = []
fixed_tracers = []
fixed_index = 0
poincare_array = np.zeros(
(tracers.x0[i_proc :: self.params.n_proc].shape[0], tracers.x0.shape[1])
)
for ix in ix0[i_proc :: self.params.n_proc]:
for iy in iy0:
# Compute Poincare index around this cell (!= 0 for potential fixed point).
diff[0, :] = np.array(
[
tracers.x1[ix, iy, tidx] - tracers.x0[ix, iy, tidx],
tracers.y1[ix, iy, tidx] - tracers.y0[ix, iy, tidx],
]
)
diff[1, :] = np.array(
[
tracers.x1[ix + 1, iy, tidx] - tracers.x0[ix + 1, iy, tidx],
tracers.y1[ix + 1, iy, tidx] - tracers.y0[ix + 1, iy, tidx],
]
)
diff[2, :] = np.array(
[
tracers.x1[ix + 1, iy + 1, tidx]
- tracers.x0[ix + 1, iy + 1, tidx],
tracers.y1[ix + 1, iy + 1, tidx]
- tracers.y0[ix + 1, iy + 1, tidx],
]
)
diff[3, :] = np.array(
[
tracers.x1[ix, iy + 1, tidx] - tracers.x0[ix, iy + 1, tidx],
tracers.y1[ix, iy + 1, tidx] - tracers.y0[ix, iy + 1, tidx],
]
)
if sum(np.sum(diff ** 2, axis=1) != 0):
diff = np.swapaxes(
np.swapaxes(diff, 0, 1) / np.sqrt(np.sum(diff ** 2, axis=1)),
0,
1,
)
poincare = self.__poincare_index(
field,
tracers.x0[ix : ix + 2, iy, tidx],
tracers.y0[ix, iy : iy + 2, tidx],
diff,
)
poincare_array[ix // self.params.n_proc, iy] = poincare
if (
abs(poincare) > 5
): # Use 5 instead of 2*pi to account for rounding errors.
# Subsample to get starting point for iteration.
nt = 2
xmin = tracers.x0[ix, iy, tidx]
ymin = tracers.y0[ix, iy, tidx]
xmax = tracers.x0[ix + 1, iy, tidx]
ymax = tracers.y0[ix, iy + 1, tidx]
xx = np.zeros((nt ** 2, 3))
tracers_part = np.zeros((nt ** 2, 5))
i1 = 0
for j1 in range(nt):
for k1 in range(nt):
xx[i1, 0] = xmin + j1 / (nt - 1.0) * (xmax - xmin)
xx[i1, 1] = ymin + k1 / (nt - 1.0) * (ymax - ymin)
xx[i1, 2] = self.params.Oz
i1 += 1
for it1 in range(nt ** 2):
# time = np.linspace(0, self.params.Lz/np.max(abs(field[2])), 10)
field_strength_z0 = vec_int(
xx[it1, :],
field,
[var.dx, var.dy, var.dz],
[var.x[0], var.y[0], var.z[0]],
[len(var.x), len(var.y), len(var.z)],
interpolation=self.params.interpolation,
)
field_strength_z0 = np.sqrt(np.sum(field_strength_z0 ** 2))
time = np.linspace(
0, 4 * self.params.Lz / field_strength_z0, 500
)
stream = Stream(field, self.params, xx=xx[it1, :], time=time)
tracers_part[it1, 0:2] = xx[it1, 0:2]
tracers_part[it1, 2:] = stream.tracers[-1, :]
min2 = 1e6
minx = xmin
miny = ymin
i1 = 0
for j1 in range(nt):
for k1 in range(nt):
diff2 = (
tracers_part[i1 + k1 * nt, 2]
- tracers_part[i1 + k1 * nt, 0]
) ** 2 + (
tracers_part[i1 + k1 * nt, 3]
- tracers_part[i1 + k1 * nt, 1]
) ** 2
if diff2 < min2:
min2 = diff2
minx = xmin + j1 / (nt - 1.0) * (xmax - xmin)
miny = ymin + k1 / (nt - 1.0) * (ymax - ymin)
# Get fixed point from this starting position using Newton's method.
point = np.array([minx, miny])
fixed_point = self.__null_point(point, var, field)
# Check if fixed point lies outside the cell.
if (
(fixed_point[0] < tracers.x0[ix, iy, tidx])
or (fixed_point[0] > tracers.x0[ix + 1, iy, tidx])
or (fixed_point[1] < tracers.y0[ix, iy, tidx])
or (fixed_point[1] > tracers.y0[ix, iy + 1, tidx])
):
fixed_point[0] = minx
fixed_point[1] = miny
# else:
fixed.append(fixed_point)
fixed_sign.append(np.sign(poincare))
fixed_index += np.sign(poincare)
# Find the streamline at the fixed point.
# time = np.linspace(0, self.params.Lz/np.max(abs(field[2])), 100)
field_strength_z0 = vec_int(
np.array([fixed_point[0], fixed_point[1], self.params.Oz]),
field,
[var.dx, var.dy, var.dz],
[var.x[0], var.y[0], var.z[0]],
[len(var.x), len(var.y), len(var.z)],
interpolation=self.params.interpolation,
)
field_strength_z0 = np.sqrt(np.sum(field_strength_z0 ** 2))
time = np.linspace(0, 4 * self.params.Lz / field_strength_z0, 500)
stream = Stream(
field,
self.params,
xx=np.array([fixed_point[0], fixed_point[1], self.params.Oz]),
time=time,
)
fixed_tracers.append(stream.tracers)
queue.put(
(i_proc, fixed, fixed_tracers, fixed_sign, fixed_index, poincare_array)
)
def find_separatrices(
self, var, field, null_point, delta=0.1, iter_max=100, ring_density=8
):
"""
find_separatrices(var, field, null_point, delta=0.1,
iter_max=100, ring_density=8)
Find the separatrices to the field 'field' with information from 'var'.
Parameters
----------
var : obj
The var object from the read_var routine.
field : string
The vector field.
null_point : obj
NullPoint object containing the magnetic null points.
delta : float
Step length for the field line tracing.
iter_max : int
Maximum iteration steps for the fiel line tracing.
ring_density : float
Density of the tracer rings.
"""
import numpy as np
from pencil.math.interpolation import vec_int
separatrices = []
connectivity = []
for null_idx in range(len(null_point.nulls)):
null = null_point.nulls[null_idx]
normal = null_point.normals[null_idx]
fan_vectors = null_point.fan_vectors[null_idx]
sign_trace = null_point.sign_trace[null_idx]
tracing = True
separatrices.append(null)
# Only trace separatrices for x-point lilke nulls.
if abs(np.linalg.det(null_point.eigen_vectors[null_idx])) < delta * 1e-8:
continue
# Create the first ring of points.
ring = []
offset = len(separatrices) - 1
for theta in np.linspace(
0, 2 * np.pi * (1 - 1.0 / ring_density), ring_density
):
ring.append(
null + self.__rotate_vector(normal, fan_vectors[0], theta) * delta
)
separatrices.append(ring[-1])
# Set the connectivity with the null point.
connectivity.append(np.array([0, len(ring)]) + offset)
# Set the connectivity within the ring.
for idx in range(ring_density - 1):
connectivity.append(np.array([idx + 1, idx + 2]) + offset)
connectivity.append(np.array([1, ring_density]) + offset)
# Trace the rings around the null.
iteration = 0
while tracing and iteration < iter_max:
ring_old = ring
# Trace field lines on ring.
point_idx = 0
for point in ring:
field_norm = vec_int(point, var, field) * sign_trace
field_norm = field_norm / np.sqrt(np.sum(field_norm ** 2))
point = point + field_norm * delta
ring[point_idx] = point
point_idx += 1
# Connectivity array between old and new ring.
connectivity_rings = np.ones((2, len(ring)), dtype="int") * range(
len(ring)
)
# Add points if distance becomes too large.
ring_new = []
ring_new.append(ring[0])
for point_idx in range(len(ring) - 1):
if self.__distance(ring[point_idx], ring[point_idx + 1]) > delta:
ring_new.append((ring[point_idx] + ring[point_idx + 1]) / 2)
connectivity_rings[1, point_idx + 1 :] += 1
ring_new.append(ring[point_idx + 1])
if self.__distance(ring[0], ring[-1]) > delta:
ring_new.append((ring[0] + ring[-1]) / 2)
ring = ring_new
# Remove points which lie outside.
ring_new = []
not_connect_to_next = []
left_shift = np.zeros(connectivity_rings.shape[1])
for point_idx in range(len(ring)):
if self.__inside_domain(ring[point_idx], var):
ring_new.append(ring[point_idx])
separatrices.append(ring[point_idx])
else:
not_connect_to_next.append(len(ring_new) - 1)
mask = connectivity_rings[1, :] == point_idx
connectivity_rings[1, mask] = -1
mask = connectivity_rings[1, :] > point_idx
left_shift += mask
connectivity_rings[1, :] -= left_shift
ring = ring_new
# Stop the tracing routine if there are no points in the ring.
if not ring:
tracing = False
continue
# Set the connectivity within the ring.
offset = len(separatrices) - len(ring_new)
for point_idx in range(len(ring_new) - 1):
if not np.any(np.array(not_connect_to_next) == point_idx):
connectivity.append(
np.array([offset + point_idx, offset + point_idx + 1])
)
if not np.any(
np.array(not_connect_to_next) == len(ring_new)
) and not np.any(np.array(not_connect_to_next) == -1):
connectivity.append(np.array([offset, offset + len(ring_new) - 1]))
# Set the connectivity between the old and new ring.
for point_old_idx in range(len(ring_old)):
if connectivity_rings[1, point_old_idx] >= 0:
connectivity_rings[0, point_old_idx] += (
len(separatrices) - len(ring_old) - len(ring)
)
connectivity_rings[1, point_old_idx] += len(separatrices) - len(
ring
)
connectivity.append(
np.array(
[
connectivity_rings[0, point_old_idx],
connectivity_rings[1, point_old_idx],
]
)
)
iteration += 1
self.separatrices = np.array(separatrices)
self.connectivity = np.array(connectivity)
def find_spines(self, var, field, null_point, delta=0.1, iter_max=100):
"""
find_spines(var, field, null_point, delta=0.1, iter_max=100)
Find the spines to the field 'field' with information from 'var'.
Parameters
----------
var : obj
The var object from the read_var routine.
field : string
The vector field.
null_point : obj
NullPoint object containing the magnetic null points.
delta : float
Step length for the field line tracing.
iter_max : int
Maximum iteration steps for the fiel line tracing.
"""
import numpy as np
from pencil.math.interpolation import vec_int
spines = []
for null_idx in range(len(null_point.nulls)):
field_sgn = -field * null_point.sign_trace[null_idx]
null = null_point.nulls[null_idx]
spine_up = []
spine_down = []
spine_up.append(null)
spine_down.append(null)
# Trace spine above the null.
iteration = 0
point = null + null_point.normals[null_idx] * delta
tracing = True
iteration = 0
while tracing and iteration < iter_max:
spine_up.append(point)
field_norm = vec_int(point, var, field_sgn)
field_norm = field_norm / np.sqrt(np.sum(field_norm ** 2))
point = point + field_norm * delta
if not self.__inside_domain(point, var):
tracing = False
iteration += 1
spines.append(np.array(spine_up))
# Trace spine below the null.
iteration = 0
point = null - null_point.normals[null_idx] * delta
tracing = True
iteration = 0
while tracing and iteration < iter_max:
spine_down.append(point)
field_norm = vec_int(point, var, field_sgn)
field_norm = field_norm / np.sqrt(np.sum(field_norm ** 2))
point = point + field_norm * delta
if not self.__inside_domain(point, var):
tracing = False
iteration += 1
spines.append(np.array(spine_down))
self.spines = np.array(spines)
def __init__(
self, field, params, xx=(0, 0, 0), time=(0, 1), metric=None, splines=None
):
"""
Trace a field starting from xx in any rectilinear coordinate system
with constant dx, dy and dz and with a given metric.
call signature:
tracer(field, params, xx=[0, 0, 0], time=[0, 1], metric=None, splines=None):
Keyword arguments:
*field*:
Vector field which is integrated over with shape [n_vars, nz, ny, nx].
Its elements are the components of the field using unnormed
unit-coordinate vectors.
*params*:
Simulation and tracer parameters.
*xx*:
Starting point of the field line integration with starting time.
*time*:
Time array for which the tracer is computed.
*metric*:
Metric function that takes a point [x, y, z] and an array
of shape [3, 3] that has the comkponents g_ij.
Use 'None' for Cartesian metric.
*splines*:
Spline interpolation functions for the tricubic interpolation.
This can speed up the calculations greatly for repeated streamline
tracing on the same data.
Accepts a list of the spline functions for the three vector components.
"""
import numpy as np
from scipy.integrate import ode
from scipy.integrate import solve_ivp
from pencil.math.interpolation import vec_int
if params.interpolation == "tricubic":
try:
import warnings
with warnings.catch_warnings():
warnings.filterwarnings("ignore", category=Warning)
from eqtools.trispline import Spline
except:
print(
"Warning: Could not import eqtools.trispline.Spline for tricubic interpolation.\n"
)
print("Warning: Fall back to trilinear.")
params.interpolation = "trilinear"
if not metric:
metric = lambda xx: np.eye(3)
dxyz = np.array([params.dx, params.dy, params.dz])
oxyz = np.array([params.Ox, params.Oy, params.Oz])
nxyz = np.array([params.nx, params.ny, params.nz])
# Redefine the derivative y for the scipy ode integrator using the given parameters.
if (params.interpolation == "mean") or (params.interpolation == "trilinear"):
odeint_func = lambda t, xx: vec_int(
xx, field, dxyz, oxyz, nxyz, params.interpolation
)
if params.interpolation == "tricubic":
x = np.linspace(params.Ox, params.Ox + params.Lx, params.nx)
y = np.linspace(params.Oy, params.Oy + params.Ly, params.ny)
z = np.linspace(params.Oz, params.Oz + params.Lz, params.nz)
if splines is None:
field_x = Spline(z, y, x, field[0, ...])
field_y = Spline(z, y, x, field[1, ...])
field_z = Spline(z, y, x, field[2, ...])
else:
field_x = splines[0]
field_y = splines[1]
field_z = splines[2]
odeint_func = lambda t, xx: self.trilinear_func(
xx, field_x, field_y, field_z, params
)
del x
del y
del z
# Set up the ode solver.
methods_ode = ["vode", "zvode", "lsoda", "dopri5", "dop853"]
methods_ivp = ["RK45", "RK23", "Radau", "BDF", "LSODA"]
if params.method in methods_ode:
solver = ode(odeint_func, jac=metric)
solver.set_initial_value(xx, time[0])
solver.set_integrator(params.method, rtol=params.rtol, atol=params.atol)
self.tracers = np.zeros([len(time), 3])
self.tracers[0, :] = xx
for i, t in enumerate(time[1:]):
self.tracers[i + 1, :] = solver.integrate(t)
if params.method in methods_ivp:
self.tracers = solve_ivp(
odeint_func,
(time[0], time[-1]),
xx,
t_eval=time,
rtol=params.rtol,
atol=params.atol,
method=params.method,
).y.T
# Remove points that lie outside the domain and interpolation on the boundary.
cut_mask = (
(
(self.tracers[:, 0] > params.Ox + params.Lx)
+ (self.tracers[:, 0] < params.Ox)
)
* (not params.periodic_x)
+ (
(self.tracers[:, 1] > params.Oy + params.Ly)
+ (self.tracers[:, 1] < params.Oy)
)
* (not params.periodic_y)
+ (
(self.tracers[:, 2] > params.Oz + params.Lz)
+ (self.tracers[:, 2] < params.Oz)
)
* (not params.periodic_z)
)
if np.sum(cut_mask) > 0:
# Find the first point that lies outside.
idx_outside = np.min(np.where(cut_mask))
# Interpolate.
p0 = self.tracers[idx_outside - 1, :]
p1 = self.tracers[idx_outside, :]
lam = np.zeros([6])
if p0[0] == p1[0]:
lam[0] = np.inf
lam[1] = np.inf
else:
lam[0] = (params.Ox + params.Lx - p0[0]) / (p1[0] - p0[0])
lam[1] = (params.Ox - p0[0]) / (p1[0] - p0[0])
if p0[1] == p1[1]:
lam[2] = np.inf
lam[3] = np.inf
else:
lam[2] = (params.Oy + params.Ly - p0[1]) / (p1[1] - p0[1])
lam[3] = (params.Oy - p0[1]) / (p1[1] - p0[1])
if p0[2] == p1[2]:
lam[4] = np.inf
lam[5] = np.inf
else:
lam[4] = (params.Oz + params.Lz - p0[2]) / (p1[2] - p0[2])
lam[5] = (params.Oz - p0[2]) / (p1[2] - p0[2])
lam_min = np.min(lam[lam >= 0])
if abs(lam_min) == np.inf:
lam_min = 0
self.tracers[idx_outside, :] = p0 + lam_min * (p1 - p0)
# We don't want to cut the interpolated point (was first point outside).
cut_mask[idx_outside] = False
cut_mask[idx_outside + 1 :] = True
# Remove outside points.
self.tracers = self.tracers[~cut_mask, :].copy()
self.params = params
self.xx = xx
self.time = time
# Compute the length of the line segments.
middle_point_list = (self.tracers[1:, :] + self.tracers[:-1, :]) / 2
diff_vectors = self.tracers[1:, :] - self.tracers[:-1, :]
self.section_l = np.zeros(diff_vectors.shape[0])
for idx, middle_point in enumerate(middle_point_list):
self.section_l[idx] = np.sqrt(
np.sum(
diff_vectors[idx]
* np.sum(metric(middle_point) * diff_vectors[idx], axis=1)
)
)
self.total_l = np.sum(self.section_l)
self.iterations = len(time)
self.section_dh = time[1:] - time[:-1]
self.total_h = time[-1] - time[0]
def __sub_tracers(self, queue, field, t_idx, i_proc, n_proc):
import numpy as np
from pencil.calc.streamlines import Stream
from pencil.math.interpolation import vec_int
# Prepare the splines for the tricubis interpolation.
if self.params.interpolation == 'tricubic':
try:
from eqtools.trispline import Spline
x = np.linspace(self.params.Ox, self.params.Ox+self.params.Lx, self.params.nx)
y = np.linspace(self.params.Oy, self.params.Oy+self.params.Ly, self.params.ny)
z = np.linspace(self.params.Oz, self.params.Oz+self.params.Lz, self.params.nz)
field_x = Spline(z, y, x, field[0, ...])
field_y = Spline(z, y, x, field[1, ...])
field_z = Spline(z, y, x, field[2, ...])
splines = np.array([field_x, field_y, field_z])
except:
splines = None
else:
splines = None
xx = np.zeros([(self.x0.shape[0]+n_proc-1-i_proc)//n_proc,
self.x0.shape[1], 3])
xx[:, :, 0] = self.x0[i_proc:self.x0.shape[0]:n_proc, :, t_idx].copy()
xx[:, :, 1] = self.y0[i_proc:self.x0.shape[0]:n_proc, :, t_idx].copy()
xx[:, :, 2] = self.z1[i_proc:self.x0.shape[0]:n_proc, :, t_idx].copy()
# Initialize the local arrays for this core.
sub_x1 = np.zeros(xx[:, :, 0].shape)
sub_y1 = np.zeros(xx[:, :, 0].shape)
sub_z1 = np.zeros(xx[:, :, 0].shape)
sub_l = np.zeros(xx[:, :, 0].shape)
sub_curly_A = np.zeros(xx[:, :, 0].shape)
sub_ee = np.zeros(xx[:, :, 0].shape)
sub_mapping = np.zeros([xx[:, :, 0].shape[0], xx[:, :, 0].shape[1], 3])
for ix in range(i_proc, self.x0.shape[0], n_proc):
for iy in range(self.x0.shape[1]):
time = np.linspace(0, 20*self.params.Lz/field[2, 0, iy, ix], 1000)
stream = Stream(field, self.params, xx=xx[int(ix/n_proc), iy, :],
time=time, splines=splines)
sub_x1[int(ix/n_proc), iy] = stream.tracers[-1, 0]
sub_y1[int(ix/n_proc), iy] = stream.tracers[-1, 1]
sub_z1[int(ix/n_proc), iy] = stream.tracers[-1, 2]
sub_l[int(ix/n_proc), iy] = stream.total_l
if self.params.int_q == 'curly_A':
for l in range(stream.total_l):
aaInt = vec_int((stream.tracers[l+1] + stream.tracers[l])/2, self.aa,
[self.params.dx, self.params.dy, self.params.dz],
[self.params.Ox, self.params.Oy, self.params.Oz],
[self.params.nx, self.params.ny, self.params.nz],
interpolation=self.params.interpolation)
sub_curly_A[int(ix/n_proc), iy] += \
np.dot(aaInt, (stream.tracers[l+1] - stream.tracers[l]))
if self.params.int_q == 'ee':
for l in range(stream.total_l):
eeInt = vec_int((stream.tracers[l+1] + stream.tracers[l])/2, self.ee,
[self.params.dx, self.params.dy, self.params.dz],
[self.params.Ox, self.params.Oy, self.params.Oz],
[self.params.nx, self.params.ny, self.params.nz],
interpolation=self.params.interpolation)
sub_ee[int(ix/n_proc), iy] += \
np.dot(eeInt, (stream.tracers[l+1] - stream.tracers[l]))
# Create the color mapping.
# if (sub_z1[int(ix/n_proc), iy] > self.params.Oz+self.params.Lz-self.params.dz*12):
if (self.x0[ix, iy, t_idx] - sub_x1[int(ix/n_proc), iy]) > 0:
if (self.y0[ix, iy, t_idx] - sub_y1[int(ix/n_proc), iy]) > 0:
sub_mapping[int(ix/n_proc), iy, :] = [0, 1, 0]
else:
sub_mapping[int(ix/n_proc), iy, :] = [1, 1, 0]
else:
if (self.y0[ix, iy, t_idx] - sub_y1[int(ix/n_proc), iy]) > 0:
sub_mapping[int(ix/n_proc), iy, :] = [0, 0, 1]
else:
sub_mapping[int(ix/n_proc), iy, :] = [1, 0, 0]
# else:
# sub_mapping[int(ix/n_proc), iy, :] = [1, 1, 1]
queue.put((i_proc, sub_x1, sub_y1, sub_z1, sub_l, sub_mapping,
sub_curly_A, sub_ee))
