def trace(map_file, psp_coord, pfss_input, retrace=False):
# Load field line
fline_file = map_file.with_suffix('.fline')
if fline_file.exists() and not retrace:
fline = load_fline(fline_file)
else:
print('Tracing field line')
# Calculate field line
output = pfsspy.pfss(pfss_input)
tracer = pfsspy.tracing.PythonTracer()
fline = tracer.trace(psp_coord, output)[0].coords
if not retrace:
fline_xyz = np.array([fline.x / u.m, fline.y / u.m, fline.z / u.m])
np.savetxt(fline_file, fline_xyz)
fline.representation_type = 'spherical'
lon = fline.lon.to_value(u.deg)
lat = fline.lat.to_value(u.deg)
r = fline.radius.to_value(u.m)
lon, lat, r = insert_nans(lon, lat, r)
fline = astropy.coordinates.SkyCoord(lon * u.deg,
lat * u.deg,
r * u.m,
frame='heliographic_carrington',
obstime=fline.obstime)
return fline
def __init__(self, nrho, rss, gong_map=sunpy.map.Map(get_gong_map())):
self._nrho = nrho
self._rss = rss
self.gong_map = gong_map
self.input = pfsspy.Input(self.gong_map, self.nrho, self.rss)
self._output = pfsspy.pfss(self.input)
self.tracer = tracing.PythonTracer()
def dipole_result(dipole_map):
nr = 10
rss = 2.5
input = pfsspy.Input(dipole_map, nr, rss)
output = pfsspy.pfss(input)
return input, output
def test_shape(zero_map):
# Test output map shapes
input, out = zero_map
nr = input.grid.nr
nphi = input.grid.nphi
ns = input.grid.ns
out = pfsspy.pfss(input)
alr, als, alp = out.al
for comp in (alr, als, alp):
assert np.all(comp == 0)
assert alr.shape == (nphi + 1, ns + 1, nr)
assert als.shape == (nphi + 1, ns, nr + 1)
assert alp.shape == (nphi, ns + 1, nr + 1)
br, bs, bp = out.bc
for comp in (br, bs, bp):
assert np.all(comp == 0)
assert br.shape == (nphi + 2, ns + 2, nr + 1)
assert bs.shape == (nphi + 2, ns + 1, nr + 2)
assert bp.shape == (nphi + 1, ns + 2, nr + 2)
br, bs, bp = out.bg
for comp in (br, bs, bp):
assert np.all(comp == 0)
assert br.shape == (nphi + 1, ns + 1, nr + 1)
assert bs.shape == (nphi + 1, ns + 1, nr + 1)
assert bp.shape == (nphi + 1, ns + 1, nr + 1)
def test_shape(zero_map):
# Test output map shapes
input, out = zero_map
nr = input.grid.nr
nphi = input.grid.nphi
ns = input.grid.ns
out = pfsspy.pfss(input)
alr, als, alp = out._al
for comp in (alr, als, alp):
assert np.all(comp == 0)
assert alr.shape == (nphi + 1, ns + 1, nr)
assert als.shape == (nphi + 1, ns, nr + 1)
assert alp.shape == (nphi, ns + 1, nr + 1)
br, bs, bp = out.bc
for comp in (br, bs, bp):
assert np.all(comp == 0)
assert br.shape == (nphi, ns, nr + 1)
assert bs.shape == (nphi, ns + 1, nr)
assert bp.shape == (nphi + 1, ns, nr)
bg = out.bg
assert np.all(bg == 0)
assert bg.shape == (nphi + 1, ns + 1, nr + 1, 3)
assert (bg[0, ...] == bg[-1, ...]).all()
def trace_to_surface(seeds, synoptic_map, rss=rss, nrho=60):
"""
seeds : astropy.coordinates.SkyCoord
Seed coordinates on the source surface.
synoptic_map : sunpy.map.GenericMap
Input synoptic magnetogram.
rss : scalar
Source surface radius.
nrho : int
Number of grid points in the radial direciton
of the PFSS model.
Returns
-------
flines : pfsspy.flines.field_lines
Traced field lines.
pfss_input : pfsspy.Input
PFSS input.
pfss_output : pfsspy.Output
PFSS output.
"""
pfss_input = pfsspy.Input(synoptic_map, nrho, rss)
print('Computing PFSS...')
pfss_output = pfsspy.pfss(pfss_input)
tracer = pfsspy.tracing.FortranTracer(max_steps=2000)
print('Tracing field lines...')
flines = tracer.trace(seeds, pfss_output)
return flines, pfss_input, pfss_output
def adapt2pfsspy(dt,
rss=2.5,
nr=60,
adapt_source="GONG",
realization="mean",
ret_magnetogram=False,
data_dir=DefaultPath):
#Input : dt (datetime object), rss (source surface radius in Rs)
YYYYMMDD = f'{dt.year}{dt.month:02d}{dt.day:02d}'
A = {"GONG": 3, "HMI": 4}.get(adapt_source)
filepath = sorted(glob.glob(f"{data_dir}/adapt40{A}*{YYYYMMDD}*.fts"))[0]
stdout.write(f"Read in {filepath}\r")
adapt_map = fits.open(filepath)
if realization == "mean": br_adapt_ = np.mean(adapt_map[0].data, axis=0)
elif isinstance(realization, int):
br_adapt_ = adapt_map[0].data[realization, :, :]
else:
raise ValueError("realization should either be 'mean' or type int ")
# Interpolate to Strumfric Grid (const lat spacing -> const cos(lat) spacing)
lat_dr = np.linspace(-90, 90, br_adapt_.shape[0])
lon_dr = np.linspace(0, 360, br_adapt_.shape[1])
clat_dr = np.sin(np.radians(lat_dr))
clat_dr_interp = np.linspace(clat_dr[0], clat_dr[-1], len(clat_dr))
br_adapt = np.zeros(br_adapt_.shape)
for ind, _ in enumerate(lon_dr):
interper = interp1d(clat_dr, br_adapt_[:, ind])
br_adapt[:, ind] = interper(clat_dr_interp)
br_adapt -= np.mean(br_adapt)
br_adapt *= 1e5 # G -> nT
if ret_magnetogram: return br_adapt
peri_input = pfsspy.Input(br_adapt, nr, rss, dtime=dt)
peri_output = pfsspy.pfss(peri_input)
return peri_output
def compute_pfss(gong_fname, dtime):
gong_map = sunpy.map.Map(gong_fname)
br = gong_map.data
header = gong_map.meta
br = br - np.mean(br)
br = np.roll(br, header['CRVAL1'] + 180, axis=1)
header['CRVAL1'] = 180
header['DATE_ORI'] = header['DATE']
header['date-obs'] = Time(dtime).isot
gong_map = sunpy.map.Map(br, header)
nrho = 60
rss = 2.5
input = pfsspy.Input(gong_map, nrho, rss)
ssfile = gong_fname.with_suffix('.ssmap')
if ssfile.exists():
ssmap = np.loadtxt(ssfile)
ssmap = sunpy.map.Map(ssmap, header)
else:
print('Calculating PFSS solution')
# Compute PFSS solution and source surface map
output = pfsspy.pfss(input)
ssdata = output.source_surface_br.data
np.savetxt(ssfile, ssdata)
ssmap = output.source_surface_br
return input, ssmap
def zero_map():
# Test a completely zero input
ns = 30
nphi = 20
nr = 10
rss = 2.5
br = np.zeros((ns, nphi))
input = pfsspy.Input(br, nr, rss)
output = pfsspy.pfss(input)
return input, output
def zero_map():
# Test a completely zero input
ns = 30
nphi = 20
nr = 10
rss = 2.5
br = np.zeros((nphi, ns))
header = pfsspy.utils.carr_cea_wcs_header(Time('1992-12-21'), br.shape)
input_map = Map((br.T, header))
input = pfsspy.Input(input_map, nr, rss)
output = pfsspy.pfss(input)
return input, output
def test_bunit(gong_map):
# Regression test to check that the output of pfss doesn't change
m = sunpy.map.Map(gong_map)
pfss_in = pfsspy.Input(m, 2, 2)
pfss_out = pfsspy.pfss(pfss_in)
assert pfss_out.bunit == u.G
pfss_out.input_map.meta['bunit'] = 'notaunit'
with pytest.warns(UserWarning, match='Could not parse unit string "notaunit"'):
assert pfss_out.bunit is None
pfss_out.input_map.meta.pop('bunit')
assert pfss_out.bunit is None
def test_pfss(gong_map):
# Regression test to check that the output of pfss doesn't change
m = sunpy.map.Map(gong_map)
# Resample to lower res for easier comparisons
m = m.resample([30, 15] * u.pix)
m = sunpy.map.Map(m.data - np.mean(m.data), m.meta)
expected = np.loadtxt(test_data / 'br_in.txt')
np.testing.assert_equal(m.data, expected)
pfss_in = pfsspy.Input(m, 50, 2)
pfss_out = pfsspy.pfss(pfss_in)
br = pfss_out.source_surface_br.data
expected = np.loadtxt(test_data / 'br_out.txt')
# atol is emperically set for tests to pass on CI
np.testing.assert_allclose(br, expected, atol=1e-13, rtol=0)
def __init__(self,
gongMap,
SSradius: float = 2.5,
nrho: int = 25,
tracer=tracing.FortranTracer()):
"""
PFSS solution to solar magnetic field is calculated on a 3D grid (phi, s, rho).
rho = ln(r), and r is the standard spherical radial coordinate.
"""
self.gongMap = gongMap
self.SSradius = SSradius
self.SSradiusRsun = SSradius * const.R_sun
self.nrho = nrho
self.input = pfsspy.Input(gongMap, nrho, SSradius)
self.output = pfsspy.pfss(self.input)
self.tracer = tracer
def dipole_map():
# Test a completely zero input
ntheta = 30
nphi = 20
nr = 10
rss = 2.5
phi = np.linspace(0, 2 * np.pi, nphi)
theta = np.linspace(-np.pi / 2, np.pi / 2, ntheta)
theta, phi = np.meshgrid(theta, phi)
def dipole_Br(r, theta):
return 2 * np.sin(theta) / r**3
br = dipole_Br(1, theta).T
input = pfsspy.Input(br, nr, rss)
output = pfsspy.pfss(input)
return input, output
def gong2pfsspy(dt,
rss=2.5,
nr=60,
ret_magnetogram=False,
data_dir=DefaultPath):
#Input : dt (datetime object), rss (source surface radius in Rs)
YYYYMMDD = f'{dt.year}{dt.month:02d}{dt.day:02d}'
gongpath = f'mrzqs{YYYYMMDD}/'
filepath = sorted(glob.glob(f'{data_dir}/mrzqs{YYYYMMDD[2:]}*.fits'))[0]
stdout.write(f"Read in {filepath}\r")
# Sunpy 1.03 error : need to fix header of gong maps to include units
am = fits.open(filepath)[0]
am.header.append(('CUNIT1', 'degree'))
am.header.append(('CUNIT2', 'degree'))
gong_map = sunpy.map.Map(am.data, am.header)
br_gong = extract_br(gong_map)
if ret_magnetogram: return br_gong
peri_input = pfsspy.Input(br_gong, nr, rss, dtime=dt)
peri_output = pfsspy.pfss(peri_input)
return peri_output
def hmi2pfsspy(dt,
rss=2.5,
nr=60,
ret_magnetogram=False,
data_dir=DefaultPath):
#Input : dt (datetime object), rss (source surface radius in Rs)
YYYYMMDD = f'{dt.year}{dt.month:02d}{dt.day:02d}'
filepath = glob.glob(
f'{data_dir}/hmi.mrdailysynframe_small_720s.{YYYYMMDD}*.fits')[0]
stdout.write(f"Read in {filepath}\r")
hmi_fits = fits.open(filepath)[0]
hmi_fits.header['CUNIT2'] = 'degree'
for card in ['HGLN_OBS', 'CRDER1', 'CRDER2', 'CSYSER1', 'CSYSER2']:
hmi_fits.header[card] = 0
hmi_map = sunpy.map.Map(hmi_fits.data, hmi_fits.header)
hmi_map.meta['CRVAL1'] = 120 + sun.L0(time=hmi_map.meta['T_OBS']).value
br_hmi = extract_br(hmi_map)
if ret_magnetogram: return br_hmi
peri_input = pfsspy.Input(br_hmi, nr, rss, dtime=dt)
peri_output = pfsspy.pfss(peri_input)
return peri_output
def derosa2pfsspy(dt,
rss=2.5,
nr=60,
ret_magnetogram=False,
data_dir=DefaultPath):
#Input : dt (datetime object), rss (source surface radius in Rs)
YYYYMMDD = f'{dt.year}{dt.month:02d}{dt.day:02d}'
dr_times = ['_000400.h5', '_060432.h5', '_120400.h5', '_180328.h5']
dr_time = dr_times[np.argmin(np.abs(dt.hour - np.array([0, 6, 12, 18])))]
filepath = f'{data_dir}/Bfield_{YYYYMMDD}{dr_time}'
f = h5py.File(filepath, 'r')
stdout.write(f"Read in {filepath}\r")
fdata = f['ssw_pfss_extrapolation']
br_dr = fdata['BR'][0, :, :, :] * 1e5
nr_dr = fdata['NR']
nlat_dr = fdata['NLAT']
nlon_dr = fdata['NLON']
r_dr = fdata['RIX'][0]
lat_dr = fdata['LAT'][0]
c_dr = np.cos(np.radians(90 - lat_dr))
lon_dr = fdata['LON'][0]
date = fdata['MODEL_DATE']
magnetogram = br_dr[0, :, :]
# Interpolate to Strumfric Grid (const lat spacing -> const cos(lat) spacing)
clat_dr = np.sin(np.radians(lat_dr))
clat_dr_interp = np.linspace(clat_dr[0], clat_dr[-1], len(clat_dr))
br_dr_interp = np.zeros(magnetogram.shape)
for ind, _ in enumerate(lon_dr):
interper = interp1d(clat_dr, magnetogram[:, ind])
br_dr_interp[:, ind] = interper(clat_dr_interp)
br_dr_interp -= np.mean(
br_dr_interp) # Remove Mean offest for pfsspy FFT based PFSS extrap
if ret_magnetogram: return (magnetogram, br_dr_interp)
peri_input = pfsspy.Input(br_dr_interp, nr, rss, dtime=dt)
peri_output = pfsspy.pfss(peri_input)
return peri_output
rss = 2.5
phi = np.linspace(0, 2 * np.pi, nphi)
theta = np.linspace(-np.pi / 2, np.pi / 2, ntheta)
theta, phi = np.meshgrid(theta, phi)
def dipole_Br(r, theta):
return 2 * np.sin(theta) / r**3
br = dipole_Br(1, theta)
br = sunpy.map.Map(br.T,
pfsspy.utils.carr_cea_wcs_header('2010-01-01', br.shape))
pfss_input = pfsspy.Input(br, nr, rss)
pfss_output = pfsspy.pfss(pfss_input)
print('Computed PFSS solution')
###############################################################################
# Trace some field lines
seed0 = np.atleast_2d(np.array([1, 1, 0]))
tracers = [pfsspy.tracing.PythonTracer(), pfsspy.tracing.FortranTracer()]
nseeds = 2**np.arange(14)
times = [[], []]
for nseed in nseeds:
print(nseed)
seeds = np.repeat(seed0, nseed, axis=0)
r, lat, lon = pfsspy.coords.cart2sph(seeds[:, 0], seeds[:, 1], seeds[:, 2])
r = r * astropy.constants.R_sun
lat = (lat - np.pi / 2) * u.rad
ax.set_ylim(0, 180)
###############################################################################
# Using the Input object, plot the input field
m = input.map
fig = plt.figure()
ax = plt.subplot(projection=m)
m.plot()
plt.colorbar()
ax.set_title('Input field')
set_axes_lims(ax)
###############################################################################
# Now calculate the PFSS solution
output = pfsspy.pfss(input)
###############################################################################
# Using the Output object we can plot the source surface field, and the
# polarity inversion line.
ss_br = output.source_surface_br
# Create the figure and axes
fig = plt.figure()
ax = plt.subplot(projection=ss_br)
# Plot the source surface map
ss_br.plot()
# Plot the polarity inversion line
ax.plot_coord(output.source_surface_pils[0])
# Plot formatting
plt.colorbar()
