Magnetic Field Flexibility

[djinpatch]

Two things that may be helpful would be the ability to add magnetic fields together and the ability to calculate the B field at a series of N,3 points all at once.

For the second point, right now it seems that if you have a list of points you must calculate in a loop like this:

points = (jnp.array([x,y,z]).T)
B_vec = []
for p in points:
    B_vec += [field.B(p)]
B_vec = jnp.array(B_vec)

It would be nice for a user to just be able to code this:

points = (jnp.array([x,y,z]).T)
B_vec = field.B(points)

For the first point, simsopt does this by defining a new class and using to to calculate the combined B field from more then one BiotSavart field. I think this could be easily done in ESSOS, where the class would have lines like this (from simsopt) which would need to be modified for Jax

class MagneticFieldSum():
.
.
.
def _B_impl(self, B):
        B[:] = np.sum([bf.B() for bf in self.Bfields], axis=0)

  def _dB_by_dX_impl(self, dB):
      dB[:] = np.sum([bf.dB_by_dX() for bf in self.Bfields], axis=0)
.
.
.

and the user could call it like this MagneticFieldSum([field1,field2]