Skip to content

Commit

Permalink
Fixed a bug from when i added openmp loops in the c++ files for curve…
Browse files Browse the repository at this point in the history 
…xyzfourier and curveplanarfourier. Might be worth going back and getting this working at some point for speed. Added option to downsample the curve-curve distance calculation. Got the Lp and SquaredMean forces and torques working, including checking the jacobians, allowing for downsampling, removing as many weak references as possible. Tried tiny speedup in the biot_savart_vjp_kernal calculation. Fixed the center function for the torque calculations. Tried edits in derivative file to no success, to try and speedup the calculations. Got the QH example fully running again with pointwise forces and net torques optimized. Remains to clean up the code and generate the new examples.
  • Loading branch information
@akaptano
akaptano committed Nov 7, 2024
1 parent 3f19636 commit b0243bf
Show file tree
Hide file tree
Showing 13 changed files with 433 additions and 359 deletions.
81 changes: 34 additions & 47 deletions examples/3_Advanced/QH_reactorscale_DA.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -142,7 +142,7 @@ def initialize_coils_QH(TEST_DIR, s):
aa = 0.05
bb = 0.05

Nx = 7
Nx = 6
Ny = Nx
Nz = Nx
# Create the initial coils:
Expand Down Expand Up @@ -264,20 +264,20 @@ def pointData_forces_torques(coils, allcoils, aprimes, bprimes, nturns_list):

LENGTH_WEIGHT = Weight(0.001)
LENGTH_TARGET = 80
LINK_WEIGHT = 1e3
LINK_WEIGHT = 1e4
CC_THRESHOLD = 0.8
CC_WEIGHT = 1e1
CS_THRESHOLD = 1.5
CS_WEIGHT = 1e2
# Weight for the Coil Coil forces term
# FORCE_WEIGHT = Weight(1e-34) # Forces are in Newtons, and typical values are ~10^5, 10^6 Newtons
# FORCE_WEIGHT2 = Weight(0.0) # Forces are in Newtons, and typical values are ~10^5, 10^6 Newtons
# TORQUE_WEIGHT = Weight(0.0) # Forces are in Newtons, and typical values are ~10^5, 10^6 Newtons
# TORQUE_WEIGHT2 = Weight(4e-27) # Forces are in Newtons, and typical values are ~10^5, 10^6 Newtons
FORCE_WEIGHT = Weight(0.0) # Forces are in Newtons, and typical values are ~10^5, 10^6 Newtons
FORCE_WEIGHT = Weight(1e-34) # Forces are in Newtons, and typical values are ~10^5, 10^6 Newtons
FORCE_WEIGHT2 = Weight(0.0) # Forces are in Newtons, and typical values are ~10^5, 10^6 Newtons
TORQUE_WEIGHT = Weight(0.0) # Forces are in Newtons, and typical values are ~10^5, 10^6 Newtons
TORQUE_WEIGHT2 = Weight(0.0) # Forces are in Newtons, and typical values are ~10^5, 10^6 Newtons
TORQUE_WEIGHT2 = Weight(1e-24) # Forces are in Newtons, and typical values are ~10^5, 10^6 Newtons
# FORCE_WEIGHT = Weight(0.0) # Forces are in Newtons, and typical values are ~10^5, 10^6 Newtons
# FORCE_WEIGHT2 = Weight(0.0) # Forces are in Newtons, and typical values are ~10^5, 10^6 Newtons
# TORQUE_WEIGHT = Weight(0.0) # Forces are in Newtons, and typical values are ~10^5, 10^6 Newtons
# TORQUE_WEIGHT2 = Weight(0.0) # Forces are in Newtons, and typical values are ~10^5, 10^6 Newtons
# Directory for output
OUT_DIR = ("./QH_reactorscale_TForder{:d}_n{:d}_p{:.2e}_c{:.2e}_lw{:.2e}_lt{:.2e}_lkw{:.2e}" + \
"_cct{:.2e}_ccw{:.2e}_cst{:.2e}_csw{:.2e}_fw{:.2e}_fww{:2e}_tw{:.2e}_tww{:2e}/").format(
Expand Down Expand Up @@ -309,8 +309,8 @@ def pointData_forces_torques(coils, allcoils, aprimes, bprimes, nturns_list):
NetTorques=coil_net_torques(coils, coils + coils_TF, regularization_rect(np.ones(len(coils)) * aa, np.ones(len(coils)) * bb), np.ones(len(coils)) * nturns)
)
# Force and Torque calculations spawn a bunch of spurious BiotSavart child objects -- erase them!
# for c in (coils + coils_TF):
# c._children = set()
for c in (coils + coils_TF):
c._children = set()

pointData = {"B_N": np.sum(btot.B().reshape((nphi, ntheta, 3)) * s.unitnormal(), axis=2)[:, :, None]}
s.to_vtk(OUT_DIR + "surf_init_DA", extra_data=pointData)
Expand Down Expand Up @@ -348,30 +348,17 @@ def pointData_forces_torques(coils, allcoils, aprimes, bprimes, nturns_list):
# interlink.
linkNum = LinkingNumber(curves + curves_TF, downsample=4)

##### Note need coils_TF + coils below!!!!!!!
# Jforce2 = sum([LpCurveForce(c, coils + coils_TF,
# regularization=regularization_rect(base_a_list[i], base_b_list[i]),
# p=2, threshold=1e8) for i, c in enumerate(base_coils + base_coils_TF)])
# Jforce2 = LpCurveForce2(coils, coils_TF, p=2, threshold=1e8)
# Jforce = sum([SquaredMeanForce(c, coils + coils_TF) for c in (base_coils + base_coils_TF)])

regularization_list = np.ones(len(coils)) * regularization_rect(aa, bb)
regularization_list2 = np.ones(len(coils_TF)) * regularization_rect(a, b)
# Jforce = MixedLpCurveForce(coils, coils_TF, regularization_list, regularization_list2) # [SquaredMeanForce2(c, coils) for c in (base_coils)]
# Jforce = MixedSquaredMeanForce(coils, coils_TF)
Jforce = sum([LpCurveForce(c, coils + coils_TF, regularization_rect(a_list[i], b_list[i]), p=4, threshold=4e5 * 100, downsample=1) for i, c in enumerate(base_coils + base_coils_TF)])
Jforce2 = sum([SquaredMeanForce(c, coils + coils_TF) for c in (base_coils + base_coils_TF)])
Jtorque = sum([LpCurveTorque(c, coils + coils_TF, regularization_rect(a_list[i], b_list[i]), p=2, threshold=4e5 * 100) for i, c in enumerate(base_coils + base_coils_TF)])
# Jtorque = sum([LpCurveTorque(c, coils + coils_TF, regularization_rect(a_list[i], b_list[i]), p=2, threshold=1e5 * 100) for i, c in enumerate(base_coils + base_coils_TF)])


Jtorque2 = sum([SquaredMeanTorque(c, coils + coils_TF) for c in (base_coils_TF)])

# Jtorque2 = sum([SquaredMeanTorque(c, coils + coils_TF) for c in (base_coils + base_coils_TF)])

# Currently, all force terms involve all the coils
all_coils = coils + coils_TF
all_base_coils = base_coils + base_coils_TF
Jforce = sum([LpCurveForce(c, all_coils, regularization_rect(a_list[i], b_list[i]), p=4, threshold=4e5 * 100, downsample=1
) for i, c in enumerate(all_base_coils)])
Jforce2 = sum([SquaredMeanForce(c, all_coils, downsample=4) for c in all_base_coils])

# Jtorque = SquaredMeanTorque2(coils, coils_TF) # [SquaredMeanForce2(c, coils) for c in (base_coils)]
# Jtorque = [SquaredMeanTorque(c, coils + coils_TF) for c in (base_coils + base_coils_TF)]
# Errors creep in when downsample = 2
Jtorque = sum([LpCurveTorque(c, all_coils, regularization_rect(a_list[i], b_list[i]), p=2, threshold=4e5 * 100, downsample=4
) for i, c in enumerate(all_base_coils)])
Jtorque2 = sum([SquaredMeanTorque(c, all_coils, downsample=1) for c in all_base_coils])

JF = Jf \
+ CC_WEIGHT * Jccdist \
Expand Down Expand Up @@ -424,10 +411,10 @@ def fun(dofs):
cc_val = CC_WEIGHT * Jccdist.J()
cs_val = CS_WEIGHT * Jcsdist.J()
link_val = LINK_WEIGHT * linkNum.J()
forces_val = FORCE_WEIGHT.value * Jforce.J()
# forces_val2 = FORCE_WEIGHT2.value * Jforce2.J()
# torques_val = TORQUE_WEIGHT.value * Jtorque.J()
# torques_val2 = TORQUE_WEIGHT2.value * Jtorque2.J()
forces_val = Jforce.J()
forces_val2 = Jforce2.J()
torques_val = Jtorque.J()
torques_val2 = Jtorque2.J()
BdotN = np.mean(np.abs(np.sum(btot.B().reshape((nphi, ntheta, 3)) * s.unitnormal(), axis=2)))
BdotN_over_B = np.mean(np.abs(np.sum(btot.B().reshape((nphi, ntheta, 3)) * s.unitnormal(), axis=2))
) / np.mean(btot.AbsB())
Expand All @@ -439,14 +426,14 @@ def fun(dofs):
valuestr += f", ccObj={cc_val:.2e}"
valuestr += f", csObj={cs_val:.2e}"
valuestr += f", Lk1Obj={link_val:.2e}"
valuestr += f", forceObj={forces_val:.2e}"
# valuestr += f", forceObj2={forces_val2:.2e}"
# valuestr += f", torqueObj={torques_val:.2e}"
# valuestr += f", torqueObj2={torques_val2:.2e}"
outstr += f", F={Jforce.J():.2e}"
# outstr += f", Fnet={Jforce2.J():.2e}"
# outstr += f", T={Jtorque.J():.2e}"
# outstr += f", Tnet={Jtorque2.J():.2e}"
valuestr += f", forceObj={FORCE_WEIGHT.value * forces_val:.2e}"
valuestr += f", forceObj2={FORCE_WEIGHT2.value * forces_val2:.2e}"
valuestr += f", torqueObj={TORQUE_WEIGHT.value * torques_val:.2e}"
valuestr += f", torqueObj2={TORQUE_WEIGHT2.value * torques_val2:.2e}"
outstr += f", F={forces_val:.2e}"
outstr += f", Fnet={forces_val2:.2e}"
outstr += f", T={torques_val:.2e}"
outstr += f", Tnet={torques_val2:.2e}"
outstr += f", C-C-Sep={Jccdist.shortest_distance():.2f}, C-S-Sep={Jcsdist2.shortest_distance():.2f}"
outstr += f", Link Number = {linkNum.J()}"
outstr += f", ║∇J║={np.linalg.norm(grad):.1e}"
Expand All @@ -456,8 +443,8 @@ def fun(dofs):
sio = io.StringIO()
sortby = SortKey.CUMULATIVE
ps = pstats.Stats(pr, stream=sio).sort_stats(sortby)
ps.print_stats(20)
print(sio.getvalue())
# ps.print_stats(20)
# print(sio.getvalue())
# for c in (coils + coils_TF):
# c._children = set()
# exit()
Expand Down
12 changes: 6 additions & 6 deletions examples/3_Advanced/QH_reactorscale_notfixed.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -332,13 +332,13 @@ def pointData_forces_torques(coils, allcoils, aprimes, bprimes, nturns_list):
# coil-coil and coil-plasma distances should be between all coils

### Jcc below removed the dipoles!
Jccdist = CurveCurveDistance(curves_TF, CC_THRESHOLD, num_basecurves=len(coils_TF))
Jccdist = CurveCurveDistance(curves_TF, CC_THRESHOLD, num_basecurves=len(coils_TF), downsample=4)
Jcsdist = CurveSurfaceDistance(curves_TF, s, CS_THRESHOLD)
# Jcsdist2 = CurveSurfaceDistance(curves_TF, s, CS_THRESHOLD)

# While the coil array is not moving around, they cannot
# interlink.
linkNum = LinkingNumber(curves_TF, downsample=4)
linkNum = LinkingNumber(curves_TF, downsample=8)

##### Note need coils_TF + coils below!!!!!!!
# Jforce2 = sum([LpCurveForce(c, coils_TF,
Expand All @@ -352,7 +352,7 @@ def pointData_forces_torques(coils, allcoils, aprimes, bprimes, nturns_list):
# Jforce = MixedLpCurveForce(coils, coils_TF, regularization_list, regularization_list2) # [SquaredMeanForce2(c, coils) for c in (base_coils)]
# Jforce = MixedSquaredMeanForce(coils, coils_TF)
# Jforce = MixedLpCurveForce(coils, coils_TF, regularization_list, regularization_list2, p=4, threshold=4e5 * 100, downsample=4)
Jforce = sum([LpCurveForce(c, coils_TF, regularization_rect(a, b), p=4, threshold=4e5 * 100, downsample=4) for i, c in enumerate(base_coils_TF)])
Jforce = sum([LpCurveForce(c, coils_TF, regularization_rect(a, b), p=4, threshold=4e5 * 100, downsample=1) for i, c in enumerate(base_coils_TF)])
# Jforce2 = sum([SquaredMeanForce(c, coils_TF) for c in (base_coils_TF)])
# Jtorque = sum([LpCurveTorque(c, coils_TF, regularization_rect(a_list[i], b_list[i]), p=2, threshold=4e5 * 100) for i, c in enumerate(base_coils_TF)])
# Jtorque = sum([LpCurveTorque(c, coils_TF, regularization_rect(a_list[i], b_list[i]), p=2, threshold=1e5 * 100) for i, c in enumerate(base_coils_TF)])
Expand Down Expand Up @@ -459,7 +459,7 @@ def fun(dofs):
# cc_val = CC_WEIGHT * Jccdist.J()
# cs_val = CS_WEIGHT * Jcsdist.J()
# link_val = LINK_WEIGHT * linkNum.J()
forces_val = FORCE_WEIGHT.value * Jforce.J()
# forces_val = FORCE_WEIGHT.value * Jforce.J()
# print(JF._children, btot._children, btot.coils[-1]._children,
# btot.coils[-1].curve._children, btot._coils[-1]._children,
# Jforce._children)
Expand All @@ -478,7 +478,7 @@ def fun(dofs):
# valuestr += f", ccObj={cc_val:.2e}"
# valuestr += f", csObj={cs_val:.2e}"
# valuestr += f", Lk1Obj={link_val:.2e}"
valuestr += f", forceObj={forces_val:.2e}"
# valuestr += f", forceObj={forces_val:.2e}"
# valuestr += f", forceObj2={forces_val2:.2e}"
# valuestr += f", torqueObj={torques_val:.2e}"
# valuestr += f", torqueObj2={torques_val2:.2e}"
Expand All @@ -494,7 +494,7 @@ def fun(dofs):
print(valuestr)
pr.disable()
sio = io.StringIO()
sortby = SortKey.CUMULATIVE
sortby = SortKey.TIME
ps = pstats.Stats(pr, stream=sio).sort_stats(sortby)
ps.print_stats(20)
print(sio.getvalue())
Expand Down
7 changes: 6 additions & 1 deletion src/simsopt/_core/derivative.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -115,11 +115,16 @@ def __add__(self, other):
x = self.data
y = other.data
z = copy_numpy_dict(x)
# for k, yk in y.items():
# if k in z:
# z[k] += yk
# else:
# z[k] = yk
for k in y:
if k in z:
z[k] += y[k]
else:
z[k] = y[k].copy()
z[k] = y[k].copy() # why copy here but not in subtract?
return Derivative(z)

def __sub__(self, other):
Expand Down
74 changes: 37 additions & 37 deletions src/simsopt/field/biotsavart.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -125,45 +125,45 @@ def B_vjp(self, v):
sopp.biot_savart_vjp_graph(points, gammas, gammadashs, currents, v,
res_gamma, res_gammadash, [], [], [])
dB_by_dcoilcurrents = self.dB_by_dcoilcurrents()
res_current = [np.sum(v * self.dB_by_dcoilcurrents()[i]) for i in range(len(dB_by_dcoilcurrents))]
res_current = [np.sum(v * dB_by_dcoilcurrents[i]) for i in range(len(dB_by_dcoilcurrents))]
return sum([coils[i].vjp(res_gamma[i], res_gammadash[i], np.asarray([res_current[i]])) for i in range(len(coils))])

def B_vjp_pure(self, v):
r"""
Assume the field was evaluated at points :math:`\mathbf{x}_i, i\in \{1, \ldots, n\}` and denote the value of the field at those points by
:math:`\{\mathbf{B}_i\}_{i=1}^n`.
These values depend on the shape of the coils, i.e. on the dofs :math:`\mathbf{c}_k` of each coil.
This function returns the vector Jacobian product of this dependency, i.e.
.. math::
\{ \sum_{i=1}^{n} \mathbf{v}_i \cdot \partial_{\mathbf{c}_k} \mathbf{B}_i \}_k.
"""
coils = self._coils
t1 = time.time()
gammas = [coil.curve.gamma() for coil in coils]
gammadashs = [coil.curve.gammadash() for coil in coils]
currents = [coil.current.get_value() for coil in coils]
res_gamma = [np.zeros_like(gamma) for gamma in gammas]
res_gammadash = [np.zeros_like(gammadash) for gammadash in gammadashs]

points = self.get_points_cart_ref()
sopp.biot_savart_vjp_graph(points, gammas, gammadashs, currents, np.array(v),
res_gamma, res_gammadash, [], [], [])
# t2 = time.time()
# print(t2 - t1)
# t1 = time.time()
# dB_by_dcoilcurrents = self.dB_by_dcoilcurrents()
# # res_current = np.sum(np.sum(v[None, :, :] * np.array(self.dB_by_dcoilcurrents()), axis=-1), axis=-1)
res_current = [jnp.sum(v * self.dB_by_dcoilcurrents()[i]) for i in range(len(coils))]
# t2 = time.time()
# print(t2 - t1)
# t1 = time.time()
# sum([coils[i].vjp(res_gamma[i], res_gammadash[i], np.asarray([res_current[i]])) for i in range(len(coils))])
# t2 = time.time()
# print(t2 - t1)
return sum([coils[i].vjp(res_gamma[i], res_gammadash[i], jnp.asarray([res_current[i]])) for i in range(len(coils))])
# def B_vjp_pure(self, v):
# r"""
# Assume the field was evaluated at points :math:`\mathbf{x}_i, i\in \{1, \ldots, n\}` and denote the value of the field at those points by
# :math:`\{\mathbf{B}_i\}_{i=1}^n`.
# These values depend on the shape of the coils, i.e. on the dofs :math:`\mathbf{c}_k` of each coil.
# This function returns the vector Jacobian product of this dependency, i.e.

# .. math::

# \{ \sum_{i=1}^{n} \mathbf{v}_i \cdot \partial_{\mathbf{c}_k} \mathbf{B}_i \}_k.

# """
# coils = self._coils
# t1 = time.time()
# gammas = [coil.curve.gamma() for coil in coils]
# gammadashs = [coil.curve.gammadash() for coil in coils]
# currents = [coil.current.get_value() for coil in coils]
# res_gamma = [np.zeros_like(gamma) for gamma in gammas]
# res_gammadash = [np.zeros_like(gammadash) for gammadash in gammadashs]

# points = self.get_points_cart_ref()
# sopp.biot_savart_vjp_graph(points, gammas, gammadashs, currents, np.array(v),
# res_gamma, res_gammadash, [], [], [])
# # t2 = time.time()
# # print(t2 - t1)
# # t1 = time.time()
# # dB_by_dcoilcurrents = self.dB_by_dcoilcurrents()
# # # res_current = np.sum(np.sum(v[None, :, :] * np.array(self.dB_by_dcoilcurrents()), axis=-1), axis=-1)
# res_current = [jnp.sum(v * self.dB_by_dcoilcurrents()[i]) for i in range(len(coils))]
# # t2 = time.time()
# # print(t2 - t1)
# # t1 = time.time()
# # sum([coils[i].vjp(res_gamma[i], res_gammadash[i], np.asarray([res_current[i]])) for i in range(len(coils))])
# # t2 = time.time()
# # print(t2 - t1)
# return sum([coils[i].vjp(res_gamma[i], res_gammadash[i], jnp.asarray([res_current[i]])) for i in range(len(coils))])

def dA_by_dcoilcurrents(self, compute_derivatives=0):
points = self.get_points_cart_ref()
Expand Down
Loading

0 comments on commit b0243bf

Please sign in to comment.