diff --git a/src/simsopt/field/magneticfield.py b/src/simsopt/field/magneticfield.py index cfe03ba68..8a647633c 100644 --- a/src/simsopt/field/magneticfield.py +++ b/src/simsopt/field/magneticfield.py @@ -92,7 +92,8 @@ def to_vtk(self, filename, nr=10, nphi=10, nz=10, rmin=1.0, rmax=2.0, zmin=-0.5, contig = np.ascontiguousarray gridToVTK(filename, X, Y, Z, pointData={"B": (contig(vals[..., 0]), contig(vals[..., 1]), contig(vals[..., 2]))}) - def to_mgrid(self, filename, nr=10, nphi=4, nz=12, rmin=1.0, rmax=2.0, zmin=-0.5, zmax=0.5, nfp=1): + def to_mgrid(self, filename, nr=10, nphi=4, nz=12, rmin=1.0, rmax=2.0, zmin=-0.5, zmax=0.5, nfp=1, + include_potential=False): """Export the field to the mgrid format for free boundary calculations. The field data is represented as a single "current group". For @@ -111,6 +112,7 @@ def to_mgrid(self, filename, nr=10, nphi=4, nz=12, rmin=1.0, rmax=2.0, zmin=-0.5 zmin: Minimum value of z for the grid. zmax: Maximum value of z for the grid. nfp: Number of field periods. + include_potential: Boolean to include the vector potential A. Defaults to false. """ rs = np.linspace(rmin, rmax, nr, endpoint=True) @@ -133,11 +135,23 @@ def to_mgrid(self, filename, nr=10, nphi=4, nz=12, rmin=1.0, rmax=2.0, zmin=-0.5 br_3 = br.reshape((nphi, nz, nr)) bp_3 = bp.reshape((nphi, nz, nr)) bz_3 = bz.reshape((nphi, nz, nr)) + + if include_potential: + A = self.A_cyl() + # shape the potential components + ar, ap, az = A.T + ar_3 = ar.reshape((nphi, nz, nr)) + ap_3 = ap.reshape((nphi, nz, nr)) + az_3 = az.reshape((nphi, nz, nr)) mgrid = MGrid(nfp=nfp, nr=nr, nz=nz, nphi=nphi, rmin=rmin, rmax=rmax, zmin=zmin, zmax=zmax) - mgrid.add_field_cylindrical(br_3, bp_3, bz_3, name='simsopt_coils') + if include_potential: + mgrid.add_field_cylindrical(br_3, bp_3, bz_3, ar=ar_3, ap=ap_3, az=az_3, name='simsopt_coils') + else: + mgrid.add_field_cylindrical(br_3, bp_3, bz_3, name='simsopt_coils') + mgrid.write(filename) diff --git a/src/simsopt/field/mgrid.py b/src/simsopt/field/mgrid.py index bb946c18b..8e7528cd2 100644 --- a/src/simsopt/field/mgrid.py +++ b/src/simsopt/field/mgrid.py @@ -69,17 +69,21 @@ def __init__(self, # fname='temp', #binary=False, self.bz_arr = [] self.bp_arr = [] - def add_field_cylindrical(self, br, bp, bz, name=None): + self.ar_arr = [] + self.az_arr = [] + self.ap_arr = [] + + def add_field_cylindrical(self, br, bp, bz, ar=None, ap=None, az=None, name=None): ''' - This function saves the vector field B. - B is defined by cylindrical components. + This function saves the vector field B, and (optionally) the vector potential A, to the Mgrid object. + B and A are provided on a tensor product grid in cylindrical components. The Mgrid array assumes B is sampled linearly first in r, then z, and last phi. Python arrays use the opposite convention such that B[0] gives a (r,z) square at const phi and B[0,0] gives a radial line and const phi and z. - It is assumed that the (br,bp,bz) inputs for this function is already in a - (nphi, nz, nr) shaped array. + It is assumed that each of the inputs ``br``, ``bp``, and ``bz`` for this function are already + arrays of shape ``(nphi, nz, nr)``. The same is true for ``ar``, ``ap``, and ``az`` if they are provided. This function may be called once for each coil group, to save sets of fields that can be scaled using EXTCUR in VMEC. @@ -88,6 +92,9 @@ def add_field_cylindrical(self, br, bp, bz, name=None): br: the radial component of B field bp: the azimuthal component of B field bz: the axial component of B field + ar: the radial component of the vector potential A + ap: the azimuthal component of the vector potential A + az: the axial component of the vector potential A name: Name of the coil group ''' @@ -95,6 +102,14 @@ def add_field_cylindrical(self, br, bp, bz, name=None): self.br_arr.append(br) self.bz_arr.append(bz) self.bp_arr.append(bp) + + # add potential + if ar is not None: + self.ar_arr.append(ar) + if ap is not None: + self.az_arr.append(az) + if az is not None: + self.ap_arr.append(ap) # add coil label if (name is None): @@ -104,9 +119,10 @@ def add_field_cylindrical(self, br, bp, bz, name=None): self.coil_names.append(label) self.n_ext_cur = self.n_ext_cur + 1 + # TO-DO: this function could check for size consistency, between different fields, and for the (nr,nphi,nz) settings of a given instance - def write(self, filename): + def write(self, filename): ''' Export class data as a netCDF binary. @@ -176,6 +192,17 @@ def write(self, filename): var_br_001[:, :, :] = self.br_arr[j] var_bz_001[:, :, :] = self.bz_arr[j] var_bp_001[:, :, :] = self.bp_arr[j] + + #If the potential value is not an empty cell, then include it + if len(self.ar_arr) > 0 : + var_ar_001 = ds.createVariable('ar'+tag, 'f8', ('phi', 'zee', 'rad')) + var_ar_001[:, :, :] = self.ar_arr[j] + if len(self.ap_arr) > 0 : + var_ap_001 = ds.createVariable('ap'+tag, 'f8', ('phi', 'zee', 'rad')) + var_ap_001[:, :, :] = self.ap_arr[j] + if len(self.az_arr) > 0 : + var_az_001 = ds.createVariable('az'+tag, 'f8', ('phi', 'zee', 'rad')) + var_az_001[:, :, :] = self.az_arr[j] @classmethod def from_file(cls, filename): @@ -203,7 +230,10 @@ def from_file(cls, filename): mgrid.filename = filename mgrid.n_ext_cur = int(f.variables['nextcur'].getValue()) coil_data = f.variables['coil_group'][:] - mgrid.coil_names = [_unpack(coil_data[j]) for j in range(mgrid.n_ext_cur)] + if len(f.variables['coil_group'].dimensions) == 2: + mgrid.coil_names = [_unpack(coil_data[j]) for j in range(mgrid.n_ext_cur)] + else: + mgrid.coil_names = [_unpack(coil_data)] mgrid.mode = f.variables['mgrid_mode'][:][0].decode() mgrid.raw_coil_current = np.array(f.variables['raw_coil_cur'][:]) @@ -212,6 +242,14 @@ def from_file(cls, filename): bp_arr = [] bz_arr = [] + ar_arr = [] + ap_arr = [] + az_arr = [] + + ar = [] + ap = [] + az = [] + nextcur = mgrid.n_ext_cur for j in range(nextcur): idx = '{:03d}'.format(j+1) @@ -228,24 +266,67 @@ def from_file(cls, filename): bp_arr.append(bp) bz_arr.append(bz) + #check for potential in mgrid file + if 'ar_'+idx in f.variables: + ar = f.variables['ar_'+idx][:] + if mgrid.mode == 'S': + ar_arr.append(ar * mgrid.raw_coil_current[j]) + else: + ar_arr.append(ar) + + if 'ap_'+idx in f.variables: + ap = f.variables['ap_'+idx][:] + if mgrid.mode == 'S': + ap_arr.append(ap * mgrid.raw_coil_current[j]) + else: + ap_arr.append(ap) + + if 'az_'+idx in f.variables: + az = f.variables['az_'+idx][:] + if mgrid.mode == 'S': + az_arr.append(az * mgrid.raw_coil_current[j]) + else: + az_arr.append(az) + mgrid.br_arr = np.array(br_arr) mgrid.bp_arr = np.array(bp_arr) mgrid.bz_arr = np.array(bz_arr) + mgrid.ar_arr = np.array(ar_arr) + mgrid.ap_arr = np.array(ap_arr) + mgrid.az_arr = np.array(az_arr) + # sum over coil groups if nextcur > 1: br = np.sum(br_arr, axis=0) bp = np.sum(bp_arr, axis=0) bz = np.sum(bz_arr, axis=0) + + if len(mgrid.ar_arr) > 0: + ar = np.sum(ar_arr, axis=0) + if len(mgrid.ap_arr) > 0: + ap = np.sum(ap_arr, axis=0) + if len(mgrid.az_arr) > 0: + az = np.sum(az_arr, axis=0) else: br = br_arr[0] bp = bp_arr[0] bz = bz_arr[0] + if len(mgrid.ar_arr) > 0: + ar = ar_arr[0] + if len(mgrid.ap_arr) > 0: + ap = ap_arr[0] + if len(mgrid.az_arr) > 0: + az = az_arr[0] mgrid.br = br mgrid.bp = bp mgrid.bz = bz + mgrid.ar = ar + mgrid.ap = ap + mgrid.az = az + mgrid.bvec = np.transpose([br, bp, bz]) return mgrid diff --git a/tests/field/test_magneticfields.py b/tests/field/test_magneticfields.py index 3985708c7..ed13a4d60 100644 --- a/tests/field/test_magneticfields.py +++ b/tests/field/test_magneticfields.py @@ -1041,14 +1041,14 @@ def test_to_vtk(self): bs = BiotSavart(coils) bs.to_vtk('/tmp/bfield') - def test_to_mgrid(self): + def subtest_to_mgrid(self, include_potential): curves, currents, ma = get_ncsx_data() nfp = 3 coils = coils_via_symmetries(curves, currents, nfp, True) bs = BiotSavart(coils) with tempfile.TemporaryDirectory() as tmpdir: filename = Path(tmpdir) / "mgrid.bfield.nc" - bs.to_mgrid(filename, nfp=nfp) + bs.to_mgrid(filename, nfp=nfp, include_potential=include_potential) # Compare the B data in the file to a separate evaluation here with netcdf_file(filename, mmap=False) as f: @@ -1068,6 +1068,13 @@ def test_to_mgrid(self): assert Br.shape == (nphi, nz, nr) assert Bphi.shape == (nphi, nz, nr) assert Bz.shape == (nphi, nz, nr) + if include_potential: + Ar = f.variables["ar_001"][()] + Aphi = f.variables["ap_001"][()] + Az = f.variables["az_001"][()] + assert Ar.shape == (nphi, nz, nr) + assert Aphi.shape == (nphi, nz, nr) + assert Az.shape == (nphi, nz, nr) r = np.linspace(rmin, rmax, nr) phi = np.linspace(0, 2 * np.pi / nfp, nphi, endpoint=False) z = np.linspace(zmin, zmax, nz) @@ -1078,6 +1085,15 @@ def test_to_mgrid(self): np.testing.assert_allclose(Br[jphi, jz, jr], bs.B_cyl()[0, 0]) np.testing.assert_allclose(Bphi[jphi, jz, jr], bs.B_cyl()[0, 1]) np.testing.assert_allclose(Bz[jphi, jz, jr], bs.B_cyl()[0, 2]) + if include_potential: + np.testing.assert_allclose(Ar[jphi, jz, jr], bs.A_cyl()[0, 0]) + np.testing.assert_allclose(Aphi[jphi, jz, jr], bs.A_cyl()[0, 1]) + np.testing.assert_allclose(Az[jphi, jz, jr], bs.A_cyl()[0, 2]) + + def test_to_mgrid(self): + for include_potential in [True, False]: + with self.subTest(include_potential=include_potential): + self.subtest_to_mgrid(include_potential) def test_poloidal_field(self): B0 = 1.1 diff --git a/tests/field/test_mgrid.py b/tests/field/test_mgrid.py index 8a913f07d..c4872a247 100644 --- a/tests/field/test_mgrid.py +++ b/tests/field/test_mgrid.py @@ -18,7 +18,7 @@ TEST_DIR = (Path(__file__).parent / ".." / "test_files").resolve() test_file = TEST_DIR / 'mgrid.pnas-qa-test-lowres-standard.nc' - +test_file2 = TEST_DIR / 'mgrid_ncsx_lowres_test.nc' class Testing(unittest.TestCase): @@ -36,6 +36,13 @@ def test_from_file(self): assert mgrid.br_arr.shape == (1, 6, 11, 11) assert mgrid.br[0, 0, 0] == -1.0633399551863771 # -0.9946816978184079 + mgrid = MGrid.from_file(test_file2) + assert mgrid.rmin == 1.0 + assert mgrid.bvec.shape == (10, 12, 4, 3) + assert mgrid.bz[1, 1, 1] == -1.012339153040808 + assert mgrid.ap[1, 1, 1] == -0.3719177477496187 + + def test_add_field_cylinder(self): N_points = 5 br = np.ones(N_points) diff --git a/tests/test_files/mgrid_ncsx_lowres_test.nc b/tests/test_files/mgrid_ncsx_lowres_test.nc new file mode 100644 index 000000000..ec40923c3 Binary files /dev/null and b/tests/test_files/mgrid_ncsx_lowres_test.nc differ