En/review eg/fields surfaces objective functions mod

[eduardolneto]

Small changes on fields.py class near_axis to work as a dependecy of custom_losses. Refactoring loss functions in objective_functions.py to meet the standard of new custom losses implementation. Adding surface-coil disstance, linking number and coils forces loss functions. Mostly changes to fields.py and objective_functions.py

[EstevaoMGomes]

Remove these changes. Removing spaces is fine, but it may cause some merge issues with other people's PRs.

[eduardolneto]

I think this is fine. The space was removed and added again, so there are virtually no chnages? This also looks the same as the file in eg/analysis branch to me.

[EstevaoMGomes]

This should not be needed...

[eduardolneto]

This was related with the way near axis is constructed. Some values are technically integers at the near-axis class level. This may not be needed now. However, if we are going to use a jax pyqsc version that may actually be necessary again because I think the original version may output some of the quantities as integers. I would keep it for now and change it later.

[EstevaoMGomes]

Then I think it should be changed in JAX PyQSC and not here. But it's ok to keep now.

[EstevaoMGomes]

I like that this was simplified, but I still need to check the losses more thoroughly. An example where they are used would be useful for that.

[eduardolneto]

@EstevaoMGomes In the examples, I kept similar local loss functions because that was the exisiting structure there. I could change them to always improt from this file. However, adding those examples to this PR would imply that I would need to add pleanty of file modifications for the examples. Which is something I thought you did not want previously.

[rogeriojorge]

1. objective_functions.py:222, 234, and 243: loss_BdotN(), loss_BdotN_only(), and loss_BdotN_only_constraint() still call field_from_dofs(...), but this PR removes field_from_dofs() earlier in the file. These functions will now raise NameError. You could make a small change of these losses to accept a field directly like the new custom-loss style.

2. I didn't test locally, but seems that objective_functions.py:81, 92, 114, 135, and 137: particle-loss functions use xyz[:,0], xyz[:,1], xyz[:,2], which selects timesteps, not coordinates, since xyz has shape (n_particles, n_times, 3). Should this be xyz[:, :, 0], xyz[:, :, 1], xyz[:, :, 2]? Otherwise jnp.diff(..., axis=1) is not differentiating the intended time series.

3. objective_functions.py:155 and 170: loss_particle_Br() and loss_particle_iota() accept a timestep argument but still pass timestep=1.e-8 to Tracing. I’d replace those with timestep=timestep, otherwise the user-provided value is ignored.

Minor things

objective_functions.py:460-505, loss_lorentz_force_coils(): block_size is in the signature, static args, and docstring, but it is not used. I’d either remove it from this function for now, or actually use it to block the mutual-field calculation. Also n_points on line 465 is unused.

new_nearaxis_from_x_and_old_nearaxis and other imports like lax might be unused and could be deleted

[eduardolneto]

I have addressed all comments. Added an exampel to test surface-distance distance-distance, linkingnumber and coisl forces. There are some incositencies with coils.py. However, these shjould only be addressed in a PR once all the current PR #40 #38 #39 #37 #36 are merged. Because it will need to use the updated coils.py file.