Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

ag/cross section fix #428

Draft
wants to merge 2 commits into
base: master
Choose a base branch
from
Draft
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
90 changes: 35 additions & 55 deletions src/simsopt/geo/surface.py
Original file line number Diff line number Diff line change
Expand Up @@ -333,19 +333,15 @@ def cross_section(self, phi, thetas=None):
single curve.
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you add a description of the arguments, i.e.

inputs
-------
phi: float, cylindrical angle between [0, 2pi] defining the cross section.
thetas: optional, Union[int, arrayLike, None]. Poloidal angle values at which to evaluate the surface on the cross section. If thetas is an integer, then the same number of evenly spaced points are used. Alternatively, thetas can be an array of poloidal angles values between [0,1]. 

"""

# phi is assumed to be between [-pi, pi], so if it does not lie on that interval
# phi is assumed to be between [0, 2*pi], so if it does not lie on that interval
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we make phi and theta follow the same convention, i.e. both in [0,1]?

# we shift it by multiples of 2pi until it does
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should we be correcting the inputs or expecting phi to be within a range? If we expect phi in [0, 1] or [0, 2pi] then we users should coerce their inputs to [0, 1] or [0, 2pi]? We would then can replace this code with assert (0 <= phi) & (phi <= 1), "phi must be in [0,1]. If we make this change, we should look to make sure we aren't breaking usage in other functions.

Alternatively, in the docstring we could specify that any value of phi is acceptable. Then this code would stay.
What do you think?

phi = phi - np.sign(phi) * np.floor(np.abs(phi) / (2 * np.pi)) * (2. * np.pi)
if phi > np.pi:
phi = phi - 2. * np.pi
if phi < -np.pi:
phi = phi + 2. * np.pi

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

if phi is the cylindrical phi for our surface, can we just return the points from gamma?

# varphi are the search intervals on which we look for the cross section in

# varphi is the search intervals on which we look for the cross section in
# at constant cylindrical phi
# The cross section is sampled at a number of points (theta_resolution) poloidally.
varphi = np.asarray([0., 0.5, 1.0])

varphi = np.asarray([0., 1.])
if thetas is None:
theta = np.asarray(self.quadpoints_theta)
elif isinstance(thetas, np.ndarray):
Expand All @@ -354,64 +350,47 @@ def cross_section(self, phi, thetas=None):
theta = np.linspace(0, 1, thetas, endpoint=False)
else:
raise NotImplementedError('Need to pass int or 1d np.array to thetas')

varphigrid, thetagrid = np.meshgrid(varphi, theta)
varphigrid = varphigrid.T
thetagrid = thetagrid.T

varphigrid, thetagrid = np.meshgrid(varphi, theta, indexing='ij')

# sample the surface at the varphi and theta points
gamma = np.zeros((varphigrid.shape[0], varphigrid.shape[1], 3))
self.gamma_lin(gamma, varphigrid.flatten(), thetagrid.flatten())

# compute the cylindrical phi coordinate of each sampled point on the surface
cyl_phi = np.arctan2(gamma[:, :, 1], gamma[:, :, 0])
cyl_phi = np.where(cyl_phi < 0, cyl_phi + 2*np.pi, cyl_phi)
cyl_phi[1, :] += 2*np.pi # second row is the same as the first row, but shifted by 2pi

varphi_left = varphigrid[0, :]
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Are we trying to determine bisection bounds here?

varphi_right = varphigrid[1, :]
cyl_phi_left = cyl_phi[0, :]
cyl_phi_right = cyl_phi[1, :]

phi = phi*np.ones(varphi_left.size)
def put_angle_between_left_and_right(angle, left_bound, right_bound):
andrewgiuliani marked this conversation as resolved.
Show resolved Hide resolved
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can you add some description. function name says angle, is that phi or varphi?

# shift angle by 2pi so that it lies between left_bound and right_bound
k = np.ceil((left_bound-angle)/(2*np.pi))
shifted_angle = angle + 2*np.pi * k
if not np.all((left_bound <= shifted_angle) & (shifted_angle <= right_bound)):
import ipdb;ipdb.set_trace()
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

remove debugging code?

raise Exception("An error occured during calculation of the cross section. \
This happens when a surface 'goes back' on itself. \
The cylindrical angle is assumed to be monotonically increasing \
with varphi, which is not the case for this surface.")
return shifted_angle

# reorder varphi, theta with respect to increasing cylindrical phi
idx = np.argsort(cyl_phi, axis=0)
cyl_phi = np.take_along_axis(cyl_phi, idx, axis=0)
varphigrid = np.take_along_axis(varphigrid, idx, axis=0)

# In case the target cylindrical angle "phi" lies above the first row or below the last row,
# we must concatenate the lower row above the top row and the top row below the lower row.
# This is allowable since the data in the matrices are periodic
cyl_phi = np.concatenate((cyl_phi[-1, :][None, :] - 2. * np.pi, cyl_phi, cyl_phi[0, :][None, :] + 2. * np.pi),
axis=0)
varphigrid = np.concatenate((varphigrid[-1, :][None, :] - 1., varphigrid, varphigrid[0, :][None, :] + 1.),
axis=0)

# ensure that varphi does not have massive jumps.
diff = varphigrid[1:] - varphigrid[:-1]
pinc = np.abs(diff + 1) < np.abs(diff)
minc = np.abs(diff - 1) < np.abs(diff)
inc = pinc.astype(int) - minc.astype(int)
prefix_sum = np.cumsum(inc, axis=0)
varphigrid[1:] = varphigrid[1:] + prefix_sum

# find the subintervals in varphi on which the desired cross section lies.
# if idx_right == 0, then the subinterval must be idx_left = 0 and idx_right = 1
idx_right = np.argmax(phi <= cyl_phi, axis=0)
idx_right = np.where(idx_right == 0, 1, idx_right)
idx_left = idx_right - 1

varphi_left = varphigrid[idx_left, np.arange(idx_left.size)]
varphi_right = varphigrid[idx_right, np.arange(idx_right.size)]
cyl_phi_left = cyl_phi[idx_left, np.arange(idx_left.size)]
cyl_phi_right = cyl_phi[idx_right, np.arange(idx_right.size)]

# this function converts varphi to cylindrical phi, ensuring that the returned angle
# lies between left_bound and right_bound.
def varphi2phi(varphi_in, left_bound, right_bound):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

describe a little here as well

# convert varphi to phi, where phi lies between left_bound and right_bound
gamma = np.zeros((varphi_in.size, 3))
self.gamma_lin(gamma, varphi_in, theta)
phi = np.arctan2(gamma[:, 1], gamma[:, 0])
pinc = (phi < left_bound).astype(int)
minc = (phi > right_bound).astype(int)
phi = phi + 2. * np.pi * (pinc - minc)
return phi
angle = np.arctan2(gamma[:, 1], gamma[:, 0])
angle = np.where(angle < 0, angle+2*np.pi, angle)
shifted_angle = put_angle_between_left_and_right(angle, left_bound, right_bound)
return shifted_angle

def bisection(phia, a, phic, c):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can we use scipy bisection or root finding?

err = 1.
while err > 1e-13:
while err > 1e-10:
b = (a + c) / 2.
phib = varphi2phi(b, phia, phic)

Expand All @@ -425,7 +404,8 @@ def bisection(phia, a, phic, c):
err = np.max(np.abs(a - c))
b = (a + c) / 2.
return b


phi = put_angle_between_left_and_right(phi, cyl_phi_left, cyl_phi_right)
# bisect cyl_phi to compute the cross section
sol = bisection(cyl_phi_left, varphi_left, cyl_phi_right, varphi_right)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

just trying to understand here: it looks like, for each theta, we are doing a bisection over varphi, so that varphi2phi(varphi) = phi. Why do we need to know cyl_phi_right and cyl_phi_left?

cross_section = np.zeros((sol.size, 3))
Expand Down
Loading