## @file m3dc1_bfield.py
#  @brief containing the problem class for pyoculus ODE solver to work on M3DC1 magnetic field
#  @author Zhisong Qu (zhisong.qu@anu.edu.au)
#

from .cylindrical_bfield import CylindricalBfield
import numpy as np

## Class that used to setup the M3DC1 bfield problem used in ODE solver.
#
class M3DC1Bfield(CylindricalBfield):
    def __init__(self, B_sim, R0, Z0):
        """! Set up the M3DC1 field for use of pyoculus
        @param B_sim the M3DC1 data generated by sim.get_bfield('B', time=sometime)
        @param R0 the R coordinate of the magnetic axis
        @param Z0 the Z coordinate of the magnetic axis
        """
        self._B = B_sim
        super().__init__(R0, Z0, 1)

    def B(self, RphiZ, *args):
        """! The magnetic field, being used by parent class CylindricalBfield
        @param RphiZ array with coordinates \f$(R, \phi, Z)\f$
        @returns \f$(B_R, B_\phi, B_Z)\f$
        """
        return np.array(self._B.evaluate((RphiZ[0], RphiZ[1], RphiZ[2])))
