More robust algorithm for finding and extending past walls

[johnomotani]

I don't think this step is the major limitation to robustness, but the function that extends contours to or past walls (PsiContour.temporaryExtend) could be improved.

hypnotoad/hypnotoad/core/equilibrium.py

Lines 1910 to 1912 in 0365640

	def temporaryExtend(
	self, *, psi, extend_lower=0, extend_upper=0, ds_lower=None, ds_upper=None
	):

)

Rather than just extrapolating from the existing grid point positions, scipy.solve_ivp() could be used to follow the direction of the flux surface to place the new points, by following a vector in the direction orthogonal to $\nabla \psi$, that is the vector [eq.f_Z, -eq.f_R], where eq is an Equilibrium object.

hypnotoad/hypnotoad/core/equilibrium.py

Lines 3975 to 3993 in 0365640

	@Equilibrium.handleMultiLocationArray
	def f_R(self, R, Z):
	"""returns the R component of the vector Grad(psi)/|Grad(psi)|**2."""
	R = numpy.clip(R, Rmin, Rmax)
	Z = numpy.clip(Z, Zmin, Zmax)
	dpsidR = self.psi_func(R, Z, dx=1, grid=False)
	dpsidZ = self.psi_func(R, Z, dy=1, grid=False)
	return dpsidR / (dpsidR**2 + dpsidZ**2)
	self.f_R = f_R.__get__(self)
	@Equilibrium.handleMultiLocationArray
	def f_Z(self, R, Z):
	"""returns the Z component of the vector Grad(psi)/|Grad(psi)|**2."""
	R = numpy.clip(R, Rmin, Rmax)
	Z = numpy.clip(Z, Zmin, Zmax)
	dpsidR = self.psi_func(R, Z, dx=1, grid=False)
	dpsidZ = self.psi_func(R, Z, dy=1, grid=False)
	return dpsidZ / (dpsidR**2 + dpsidZ**2)

This would be similar to

hypnotoad/hypnotoad/core/mesh.py

Lines 3291 to 3449 in 0365640

	def followPerpendicular(
	i,
	p0,
	psi0,
	*,
	f_R,
	f_Z,
	psivals,
	rtol=2.0e-8,
	atol=1.0e-8,
	maxits: int = 1000,
	recover: bool = False,
	**kwargs,
	):
	"""Follow a line perpendicular to Bp from point p0 until psi_target is reached.
	# Arguments
	maxits : int
	Maximum number of iterations to be taken. If exceeded then the range
	of psi will be truncated.
	recover : bool
	Recover from failures by changing the psivals of
	the grid points. This will result in an incorrect grid, but
	is useful when adjusting settings.
	"""
	if i is not None:
	print(f"Following perpendicular: {i + 1}", end="\r", flush=True)
	# psi0 might be in somewhere in the range of psivals, rather than at one end
	if min(psivals) < psi0 < max(psivals):
	# Integrate in each direction, then put together
	# Partition into left and right halves
	if psivals[0] < psi0:
	left = [psi for psi in psivals if psi < psi0]
	right = [psi for psi in psivals if psi >= psi0]
	else:
	left = [psi for psi in psivals if psi >= psi0]
	right = [psi for psi in psivals if psi < psi0]
	return followPerpendicular(
	None,
	p0,
	psi0,
	f_R=f_R,
	f_Z=f_Z,
	psivals=left[::-1],
	rtol=rtol,
	atol=atol,
	maxits=maxits,
	recover=recover,
	)[::-1] + followPerpendicular(
	None,
	p0,
	psi0,
	f_R=f_R,
	f_Z=f_Z,
	psivals=right,
	rtol=rtol,
	atol=atol,
	maxits=maxits,
	recover=recover,
	)
	if abs(psivals[-1] - psi0) < abs(psivals[0] - psi0):
	# Closer at the end than the start -> Reverse
	return followPerpendicular(
	None,
	p0,
	psi0,
	f_R=f_R,
	f_Z=f_Z,
	psivals=psivals[::-1],
	rtol=rtol,
	atol=atol,
	maxits=maxits,
	recover=recover,
	)[::-1]
	psivals = psivals.copy()
	class MaxIterException(Exception):
	def __init__(self, maxits, psi):
	super().__init__(
	f"Max iterations {maxits} reached at psi = {psi}\n"
	" To recover from this and generate a grid anyway, "
	"set follow_perpendicular_recover to True.\n"
	" The resulting grid will not align to flux surfaces."
	)
	self.psi = psi
	call_counter = 0
	def f(psi, x):
	# Implement a maximum number of iterations because this is not provided by
	# the solve_ivp API.
	nonlocal call_counter
	call_counter += 1
	if call_counter >= maxits:
	raise MaxIterException(maxits, psi)
	return (f_R(x[0], x[1]), f_Z(x[0], x[1]))
	psirange = (psi0, psivals[-1])
	# make sure rounding errors do not cause exception:
	if psirange[1] - psirange[0] > 0:
	# psi increasing in this interval
	if psivals[0] < psirange[0] and psirange[0] - psivals[0] < 1.0e-15 * numpy.abs(
	psirange[0]
	):
	# rounding error present, reset psivals[0]
	psivals[0] = psirange[0]
	else:
	# psi decreasing in this interval
	if psivals[0] > psirange[0] and psivals[0] - psirange[0] < 1.0e-15 * numpy.abs(
	psirange[0]
	):
	# rounding error present, reset psivals[0]
	psivals[0] = psirange[0]
	try:
	solution = solve_ivp(
	f,
	psirange,
	tuple(p0),
	t_eval=psivals,
	rtol=rtol,
	atol=atol,
	vectorized=True,
	)
	except MaxIterException as e:
	print(f"followPerpendicular failed at psi = {e.psi}")
	if recover:
	print(
	" WARNING: Recovering by changing psi values.\n"
	" The resulting grid will not align to flux surfaces."
	)
	new_psivals = numpy.linspace(psi0, e.psi, len(psivals) + 1, endpoint=False)[
	1:
	]
	return followPerpendicular(
	None,
	p0,
	psi0,
	f_R=f_R,
	f_Z=f_Z,
	psivals=new_psivals,
	rtol=rtol,
	atol=atol,
	maxits=maxits,
	recover=recover,
	)
	else:
	raise e
	except ValueError:
	print(psirange, psivals)
	raise
	return [Point2D(*p) for p in solution.y.T]

but where that follows perpendicular to flux surfaces, this would follow along them.

[johnomotani]

See also #146.