def test_crop_rotated_celestial(ndcube_4d_ln_lt_l_t):
# This is a regression test for a highly rotated image where all 4 corners
# of the spatial ROI have to be used.
header = dedent("""\
WCSAXES = 2 / Number of coordinate axes
CRPIX1 = 2053.459961 / Pixel coordinate of reference point
CRPIX2 = 2047.880005 / Pixel coordinate of reference point
PC1_1 = 0.70734471922412 / Coordinate transformation matrix element
PC1_2 = 0.70686876305701 / Coordinate transformation matrix element
PC2_1 = -0.70686876305701 / Coordinate transformation matrix element
PC2_2 = 0.70734471922412 / Coordinate transformation matrix element
CDELT1 = 0.00016652472222222 / [deg] Coordinate increment at reference point
CDELT2 = 0.00016652472222222 / [deg] Coordinate increment at reference point
CUNIT1 = 'deg' / Units of coordinate increment and value
CUNIT2 = 'deg' / Units of coordinate increment and value
CTYPE1 = 'HPLN-TAN' / Coordinate type codegnomonic projection
CTYPE2 = 'HPLT-TAN' / Coordinate type codegnomonic projection
CRVAL1 = 0.0 / [deg] Coordinate value at reference point
CRVAL2 = 0.0 / [deg] Coordinate value at reference point
LONPOLE = 180.0 / [deg] Native longitude of celestial pole
LATPOLE = 0.0 / [deg] Native latitude of celestial pole
MJDREF = 0.0 / [d] MJD of fiducial time
DATE-OBS= '2014-04-09T06:00:12.970' / ISO-8601 time of observation
MJD-OBS = 56756.250150116 / [d] MJD of observation
RSUN_REF= 696000000.0 / [m] Solar radius
DSUN_OBS= 149860273889.04 / [m] Distance from centre of Sun to observer
HGLN_OBS= -0.0058904803279347 / [deg] Stonyhurst heliographic lng of observer
HGLT_OBS= -6.0489216362492 / [deg] Heliographic latitude of observer
""")
wcs = WCS(fits.Header.fromstring(header, sep="\n"))
data = np.zeros((4096, 4096))
cube = NDCube(data, wcs=wcs)
bottom_left = SkyCoord(-100,
-100,
unit=u.arcsec,
frame=wcs_to_celestial_frame(wcs))
bottom_right = SkyCoord(600,
-100,
unit=u.arcsec,
frame=wcs_to_celestial_frame(wcs))
top_left = SkyCoord(-100,
600,
unit=u.arcsec,
frame=wcs_to_celestial_frame(wcs))
top_right = SkyCoord(600,
600,
unit=u.arcsec,
frame=wcs_to_celestial_frame(wcs))
small = cube.crop(bottom_left, bottom_right, top_left, top_right)
assert small.data.shape == (1652, 1652)
def ndcube_3d_rotated(wcs_3d_ln_lt_t_rotated, simple_extra_coords_3d):
data_rotated = np.array([[[1, 2, 3, 4, 6], [2, 4, 5, 3, 1],
[0, -1, 2, 4, 2], [3, 5, 1, 2, 0]],
[[2, 4, 5, 1, 3], [1, 5, 2, 2, 4],
[2, 3, 4, 0, 5], [0, 1, 2, 3, 4]]])
mask_rotated = data_rotated >= 0
cube = NDCube(
data_rotated,
wcs_3d_ln_lt_t_rotated,
mask=mask_rotated,
uncertainty=data_rotated,
)
cube._extra_coords = simple_extra_coords_3d
return cube
def ndcube_3d_ln_lt_l_ec_time(wcs_3d_l_lt_ln, time_and_simple_extra_coords_2d):
shape = (2, 3, 4)
wcs_3d_l_lt_ln.array_shape = shape
data = data_nd(shape)
mask = data > 0
cube = NDCube(
data,
wcs_3d_l_lt_ln,
mask=mask,
uncertainty=data,
)
cube._extra_coords = time_and_simple_extra_coords_2d
cube._extra_coords._ndcube = cube
return cube
def ndcube_3d_ln_lt_l_ec_all_axes(wcs_3d_l_lt_ln, extra_coords_3d):
shape = (2, 3, 4)
wcs_3d_l_lt_ln.array_shape = shape
data = data_nd(shape)
mask = data > 0
cube = NDCube(
data,
wcs_3d_l_lt_ln,
mask=mask,
uncertainty=data,
)
cube._extra_coords = extra_coords_3d
cube._extra_coords._ndcube = cube
return cube
def gen_ndcube_3d_l_ln_lt_ectime(wcs_3d_lt_ln_l,
time_axis,
time_base,
global_coords=None):
shape = (10, 5, 8)
wcs_3d_lt_ln_l.array_shape = shape
data_cube = data_nd(shape)
mask = data_cube < 0
extra_coords = time_extra_coords(shape, time_axis, time_base)
cube = NDCube(data_cube, wcs_3d_lt_ln_l, mask=mask, uncertainty=data_cube)
cube._extra_coords = extra_coords
if global_coords:
cube._global_coords = global_coords
return cube
def total_intensity(self):
"""
The intensity summed over an entire spectral window.
"""
data = np.sum(self.data, axis=0)
wcs = self.wcs.dropaxis(2)
return NDCube(data, wcs)
def ndcube_4d_unit_uncertainty(wcs_4d_t_l_lt_ln):
shape = (5, 8, 10, 12)
data_cube = data_nd(shape)
uncertainty = np.sqrt(data_cube)
return NDCube(data_cube,
wcs=wcs_4d_t_l_lt_ln,
unit=u.J,
uncertainty=uncertainty)
def fromFile(cls, files: str) -> 'ViewerController':
cubes = [NDCube(getdata(f), WCS(getheader(f))) for f in files]
if len(cubes) == 1:
cube = cubes[0]
else:
cube = NDCubeSequence(cubes)
model = NDCubeModel(cube, files)
return cls(model)
def test_add_coord_after_create(time_lut):
ndc = NDCube(np.random.random((10, 10)), wcs=WCS(naxis=2))
assert isinstance(ndc.extra_coords, ExtraCoords)
ndc.extra_coords.add("time", 0, time_lut)
assert len(ndc.extra_coords._lookup_tables) == 1
assert ndc.extra_coords["time"]._lookup_tables == ndc.extra_coords._lookup_tables
def test_wcs_type_after_init(ndcube_3d_ln_lt_l, wcs_3d_l_lt_ln):
# Generate a low level WCS
slices = np.s_[:, :, 0]
low_level_wcs = SlicedLowLevelWCS(wcs_3d_l_lt_ln, slices)
# Generate an NDCube using the low level WCS
cube = NDCube(ndcube_3d_ln_lt_l.data[slices], low_level_wcs)
# Check the WCS has been converted to high level but NDCube init.
assert isinstance(cube.wcs, BaseHighLevelWCS)
def ndcube_4d_mask(wcs_4d_t_l_lt_ln):
shape = (5, 8, 10, 12)
data_cube = data_nd(shape)
uncertainty = np.sqrt(data_cube)
mask = data_cube % 2
return NDCube(data_cube,
wcs=wcs_4d_t_l_lt_ln,
uncertainty=uncertainty,
mask=mask)
def test_initialize_from_ndcube(ndcube_3d_l_ln_lt_ectime):
cube = ndcube_3d_l_ln_lt_ectime
cube.global_coords.add('distance', 'pos.distance', 1 * u.m)
cube2 = NDCube(cube)
assert cube.global_coords is cube2.global_coords
assert cube.extra_coords is cube2.extra_coords
cube3 = NDCube(cube, copy=True)
ec = cube.extra_coords
ec3 = cube3.extra_coords
assert cube.global_coords == cube3.global_coords
assert cube.global_coords is not cube3.global_coords
assert ec.keys() == ec3.keys()
assert ec.mapping == ec3.mapping
assert np.allclose(ec.wcs.pixel_to_world_values(1), ec3.wcs.pixel_to_world_values(1))
assert ec is not ec3
def ndcube_3d_ln_lt_l(wcs_3d_l_lt_ln, simple_extra_coords_3d):
shape = (2, 3, 4)
wcs_3d_l_lt_ln.array_shape = shape
data = data_nd(shape)
mask = data > 0
return NDCube(data,
wcs_3d_l_lt_ln,
mask=mask,
uncertainty=data,
extra_coords=simple_extra_coords_3d)
def test_collection_update_collecton_input():
orig_collection = NDCollection([("cube0", cube0), ("cube1", cube1)],
aligned_axes[:2])
cube1_alt = NDCube(data1 * 2, input_wcs1)
new_collection = NDCollection([("cube1", cube1_alt), ("cube2", cube2)],
aligned_axes[1:])
orig_collection.update(new_collection)
expected = NDCollection([("cube0", cube0), ("cube1", cube1_alt),
("cube2", cube2)], aligned_axes)
helpers.assert_collections_equal(orig_collection, expected)
def test_combined_wcs(time_lut):
ndc = NDCube(np.random.random((10, 10)), wcs=WCS(naxis=2))
assert isinstance(ndc.extra_coords, ExtraCoords)
ndc.extra_coords.add_coordinate("time", 0, time_lut)
cwcs = ndc.combined_wcs
assert cwcs.world_n_dim == 3
assert cwcs.pixel_n_dim == 2
world = cwcs.pixel_to_world(0, 0)
assert u.allclose(world[:2], (1, 1) * u.one)
assert world[2] == Time("2011-01-01T00:00:00")
def _convert_iris_sequence(sequence, new_unit):
"""Converts data and uncertainty in an IRISSpectrogramSequence between units.
Parameters
----------
sequence: `NDCubeSequence`, `SpectrogramSequence` or `IRISSpectrogramSequence`
Sequence whose constituent NDCubes are be converted to new units.
new_unit: `astropy.units.Unit` or `str`
Unit to which the data is to be converted.
Returns
-------
converted_data_list: `list` of `NDCube`s.
List of NDCubes with data and uncertainty attributes converted to new_unit.
"""
# Define empty list to hold NDCubes with converted data and uncertainty.
converted_data_list = []
# Cycle through each NDCube, convert data and uncertainty to new
# units, and append to list.
for i, cube in enumerate(sequence.data):
# Determine what type of DN unit is needed based on detector type.
detector_type = _get_detector_type(cube.meta)
if new_unit == "DN":
new_unit = DN_UNIT[detector_type]
# If NDCube is already in new unit, add NDCube as is to list.
if cube.unit is new_unit or cube.unit is new_unit / u.s:
converted_data_list.append(cube)
# Else convert data and uncertainty to new unit.
if cube.unit != new_unit or cube.unit != new_unit / u.s:
# During calculations, the time component due to exposure
# time correction, if it has been applied, is ignored.
# Check here whether the time correction is present in the
# original unit so that is carried through to new unit.
if u.s not in (cube.unit.decompose() * u.s).bases:
new_unit_time_accounted = new_unit / u.s
else:
new_unit_time_accounted = new_unit
# Convert data and uncertainty to new unit.
data = (cube.data * cube.unit).to(new_unit).value
uncertainty = (cube.uncertainty.array *
cube.unit).to(new_unit).value
# Append new instance of NDCube in new unit to list.
converted_data_list.append(
NDCube(data,
wcs=cube.wcs,
meta=cube.meta,
mask=cube.mask,
unit=new_unit_time_accounted,
uncertainty=uncertainty,
extra_coords=_extra_coords_to_input_format(
cube._extra_coords)))
return converted_data_list
def _apply_or_undo_exposure_time_correction(sequence, correction_function):
"""Applies or undoes exposure time correction to a sequence of NDCubes.
Correction is applied (or undone) to both data and uncertainty attributes of NDCubes.
Parameters
----------
sequence: `NDCubeSequence`, `SpectrogramSequence` or `IRISSpectrogramSequence`
Sequence whose constituent NDCubes are be converted to new units.
NDCubes with sequence must have an 'exposure time' entry in its extra
coords attribute.
correction_function: function
Function applying or undoing exposure time correction.
Returns
-------
converted_data_list: `list` of `NDCube`s.
List of NDCubes with data and uncertainty corrected (or uncorrected)
for exposure time.
"""
converted_data_list = []
for i, cube in enumerate(sequence.data):
if u.s not in cube.unit.decompose().bases:
exposure_time_s = cube._extra_coords["exposure time"]["value"].to(
u.s).value
if len(cube.dimensions.shape) == 1:
pass
elif len(cube.dimensions.shape) == 2:
exposure_time_s = exposure_time_s[:, np.newaxis]
elif len(cube.dimensions.shape) == 3:
exposure_time_s = exposure_time_s[:, np.newaxis, np.newaxis]
else:
raise ValueError(
"NDCube dimensions must be 2 or 3. Dimensions={0}".format(
len(cube.dimensions.shape)))
data = correction_function(cube.data, exposure_time_s)
uncertainty = correction_function(cube.uncertainty.array,
exposure_time_s)
converted_data_list.append(
NDCube(data,
wcs=cube.wcs,
meta=cube.meta,
mask=cube.mask,
unit=cube.unit / u.s,
uncertainty=uncertainty,
extra_coords=_extra_coords_to_input_format(
cube._extra_coords)))
else:
converted_data_list.append(cube)
return converted_data_list
def test_dropped_dimension_reordering():
data = np.ones((3, 4, 5))
wcs_input_dict = {
'CTYPE1': 'WAVE ', 'CUNIT1': 'Angstrom', 'CDELT1': 0.2, 'CRPIX1': 0, 'CRVAL1': 10, 'NAXIS1': 5,
'CTYPE2': 'HPLT-TAN', 'CUNIT2': 'deg', 'CDELT2': 0.5, 'CRPIX2': 2, 'CRVAL2': 0.5, 'NAXIS2': 4,
'CTYPE3': 'HPLN-TAN', 'CUNIT3': 'deg', 'CDELT3': 0.4, 'CRPIX3': 2, 'CRVAL3': 1, 'NAXIS3': 3}
input_wcs = WCS(wcs_input_dict)
base_time = Time('2000-01-01', format='fits', scale='utc')
timestamps = Time([base_time + TimeDelta(60 * i, format='sec') for i in range(data.shape[0])])
my_cube = NDCube(data, input_wcs)
my_cube.extra_coords.add('time', (0,), timestamps)
# If the argument to extra_coords.add is array index then it should end up
# in the first element of array_axis_physical_types
assert "time" in my_cube.array_axis_physical_types[0]
# When we slice out the dimension with the extra coord in it should go away.
assert "time" not in my_cube[0].array_axis_physical_types[0]
def ndcube_4d_ln_lt_l_t(wcs_4d_t_l_lt_ln):
shape = (5, 8, 10, 12)
wcs_4d_t_l_lt_ln.array_shape = shape
data_cube = data_nd(shape)
return NDCube(data_cube, wcs=wcs_4d_t_l_lt_ln)
}
wm = WCS(header=hm, naxis=3)
data = np.array([[[1, 2, 3, 4], [2, 4, 5, 3], [0, -1, 2, 3]],
[[2, 4, 5, 1], [10, 5, 2, 2], [10, 3, 3, 0]]])
uncertainty = np.sqrt(data)
mask_cube = data < 0
cube = NDCube(data,
wt,
mask=mask_cube,
uncertainty=uncertainty,
missing_axis=[False, False, False, True],
extra_coords=[
('time', 0, u.Quantity(range(data.shape[0]), unit=u.s)),
('hello', 1, u.Quantity(range(data.shape[1]), unit=u.W)),
('bye', 2, u.Quantity(range(data.shape[2]), unit=u.m)),
('another time', 2,
np.array([
datetime.datetime(2000, 1, 1) +
datetime.timedelta(minutes=i)
for i in range(data.shape[2])
])), ('array coord', 2, np.arange(100, 100 + data.shape[2]))
])
cube_unit = NDCube(
data,
wt,
mask=mask_cube,
unit=u.J,
uncertainty=uncertainty,
missing_axis=[False, False, False, True],
'CUNIT1': 'deg',
'CDELT1': 0.5,
'CRPIX1': 2,
'CRVAL1': 0.5,
'NAXIS1': 4,
'CTYPE2': 'HPLN-TAN',
'CUNIT2': 'deg',
'CDELT2': 0.4,
'CRPIX2': 2,
'CRVAL2': 1,
'NAXIS2': 3
}
input_wcs1 = astropy.wcs.WCS(wcs_input_dict1)
# Define cubes.
cube0 = NDCube(data0, input_wcs)
cube1 = NDCube(data1, input_wcs1)
cube2 = NDCube(data2, input_wcs)
# Define sequences.
sequence02 = NDCubeSequence([cube0, cube2])
sequence20 = NDCubeSequence([cube2, cube0])
# Define collections
aligned_axes = ((1, 2), (2, 0), (1, 2))
keys = ("cube0", "cube1", "cube2")
cube_collection = NDCollection([("cube0", cube0), ("cube1", cube1),
("cube2", cube2)], aligned_axes)
seq_collection = NDCollection([("seq0", sequence02), ("seq1", sequence20)],
aligned_axes="all")
'NAXIS2': 3,
'CTYPE3': 'HPLN-TAN',
'CUNIT3': 'deg',
'CDELT3': 0.4,
'CRPIX3': 2,
'CRVAL3': 1,
'NAXIS3': 2,
}
wt = WCS(header=ht, naxis=3)
wm = WCS(header=hm, naxis=3)
cube1 = NDCube(data,
wt,
missing_axes=[False, False, False, True],
extra_coords=[
('pix', 0, u.Quantity(range(data.shape[0]), unit=u.pix)),
('distance', None, u.Quantity(0, unit=u.cm)),
('time', None, datetime.datetime(2000, 1, 1, 0, 0))
])
cube2 = NDCube(data,
wm,
extra_coords=[
('pix', 0,
u.Quantity(np.arange(1, data.shape[0] + 1), unit=u.pix) +
cube1.extra_coords['pix']['value'][-1]),
('distance', None, u.Quantity(1, unit=u.cm)),
('time', None, datetime.datetime(2000, 1, 1, 0, 1))
])
cube3 = NDCube(data2,
'CRVAL2': 0.5,
'NAXIS2': 3,
'CTYPE3': 'HPLN-TAN',
'CUNIT3': 'deg',
'CDELT3': 0.4,
'CRPIX3': 2,
'CRVAL3': 1,
'NAXIS3': 2
}
wm = WCS(header=hm, naxis=3)
cube1 = NDCube(data,
wt,
missing_axes=[False, False, False, True],
extra_coords=[
('pix', 0, u.Quantity(range(data.shape[0]), unit=u.pix)),
('hi', 1, u.Quantity(range(data.shape[1]), unit=u.s)),
('distance', None, u.Quantity(0, unit=u.cm)),
('time', None, datetime.datetime(2000, 1, 1, 0, 0))
])
cube1_with_unit = NDCube(
data,
wt,
missing_axes=[False, False, False, True],
unit=u.km,
extra_coords=[('pix', 0, u.Quantity(range(data.shape[0]), unit=u.pix)),
('hi', 1, u.Quantity(range(data.shape[1]), unit=u.s)),
('distance', None, u.Quantity(0, unit=u.cm)),
('time', None, datetime.datetime(2000, 1, 1, 0, 0))])
def ndcube_1d_l(wcs_1d_l):
shape = (10, )
data_cube = data_nd(shape)
return NDCube(data_cube, wcs=wcs_1d_l)
def ndcube_2d_ln_lt(wcs_2d_lt_ln):
shape = (10, 12)
data_cube = data_nd(shape)
return NDCube(data_cube, wcs=wcs_2d_lt_ln)
def stack_spectrogram_sequence(cube_sequence, memmap=True, reproject=False):
"""
Given a sequence of IRIS rasters stack them into a single `ndcube.NDCube`.
.. warning::
This is intended to be used for plotting only, it will not preserve the
flux in the image.
Parameters
----------
cube_sequence : `irispy.spectrogram.IRISSpectrogramCubeSequence`
The input arrays to regrid.
memmap : `bool`
Use a temporary file to store the re-gridded data in rather than in memory.
Returns
-------
`ndcube.NDCube`: A 4D cube with a new time dimension
"""
if len(cube_sequence.data) == 1:
raise ValueError("No point doing this to one raster")
if memmap:
if not isinstance(memmap, Path):
memmap = tempfile.TemporaryFile()
target_wcs = cube_sequence[0].wcs
target_shape = cube_sequence[0].data.shape
cube_shape = tuple([len(cube_sequence.data)] + list(target_shape))
memmap = np.memmap(
memmap, cube_sequence[0].data.dtype, "w+",
shape=cube_shape) if memmap else np.empty(cube_shape)
times = [cube_sequence[0].extra_coords['time']['value'][0]]
memmap[0] = cube_sequence[0].data
for i, cube in enumerate(cube_sequence.data[1:]):
if not reproject:
memmap[i + 1] = cube_sequence[i + 1].data
else:
reproject_interp((cube.data, cube.wcs),
target_wcs,
shape_out=target_shape,
independent_celestial_slices=False,
order=0,
return_footprint=False,
output_array=memmap[i + 1])
times.append(cube.extra_coords['time']['value'][0])
times = Time(times)
dts = times[1:] - times[:-1]
if u.allclose(dts[0].to(u.s), dts.to(u.s), atol=0.5 * u.s):
dt = dts[0]
else:
raise ValueError("Can't handle tabular wcs")
out_wcs = WCS(naxis=4)
out_wcs.wcs.crpix = list(target_wcs.wcs.crpix) + [0]
out_wcs.wcs.crval = list(target_wcs.wcs.crval) + [0]
out_wcs.wcs.cdelt = list(target_wcs.wcs.cdelt) + [dt.to(u.s).value]
out_wcs.wcs.ctype = list(target_wcs.wcs.ctype) + ['TIME']
out_wcs.wcs.cunit = list(target_wcs.wcs.cunit) + ['s']
pc = np.identity(4)
pc[:3, :3] = target_wcs.wcs.pc
out_wcs.wcs.pc = pc
return NDCube(memmap, out_wcs)
mask_cubem = data > 0
mask_cube = data >= 0
uncertaintym = data
uncertainty = np.sqrt(data)
mask_disordered = data_disordered > 0
uncertainty_disordered = data_disordered
mask_ordered = data_ordered > 0
uncertainty_ordered = data_ordered
cubem = NDCube(
data,
wm,
mask=mask_cubem,
uncertainty=uncertaintym,
extra_coords=[('time', 0, u.Quantity(range(data.shape[0]), unit=u.pix)),
('hello', 1, u.Quantity(range(data.shape[1]), unit=u.pix)),
('bye', 2, u.Quantity(range(data.shape[2]), unit=u.pix))])
cube_disordered_inputs = (
data_disordered, w_disordered, mask_disordered, uncertainty_disordered,
[('spam', 0, u.Quantity(range(data_disordered.shape[0]), unit=u.pix)),
('hello', 1, u.Quantity(range(data_disordered.shape[1]), unit=u.pix)),
('bye', 2, u.Quantity(range(data_disordered.shape[2]), unit=u.pix))])
cube_disordered = NDCube(cube_disordered_inputs[0], cube_disordered_inputs[1],
mask=cube_disordered_inputs[2], uncertainty=cube_disordered_inputs[3],
extra_coords=cube_disordered_inputs[4])
cube_ordered = NDCubeOrdered(
data_ordered,
def ndcube_4d_extra_coords(wcs_4d_t_l_lt_ln, simple_extra_coords_3d):
shape = (5, 8, 10, 12)
data_cube = data_nd(shape)
return NDCube(data_cube,
wcs=wcs_4d_t_l_lt_ln,
extra_coords=simple_extra_coords_3d)
'CDELT2': 0.5,
'CRPIX2': 2,
'CRVAL2': 0.5,
'NAXIS2': 3,
'CTYPE3': 'HPLN-TAN',
'CUNIT3': 'deg',
'CDELT3': 0.4,
'CRPIX3': 2,
'CRVAL3': 1,
'NAXIS3': 2,
}
wt = WCS(header=ht, naxis=3)
wm = WCS(header=hm, naxis=3)
cube1 = NDCube(data, wt, missing_axis=[False, False, False, True])
cube2 = NDCube(data, wm)
cube3 = NDCube(data2, wt, missing_axis=[False, False, False, True])
cube4 = NDCube(data2, wm)
seq = NDCubeSequence([cube1, cube2, cube3, cube4], common_axis=0)
seq1 = NDCubeSequence([cube1, cube2, cube3, cube4])
map1 = cube2[:, :, 0].to_sunpy()
map2 = cube2[:, :, 1].to_sunpy()
map3 = cube2[:, :, 2].to_sunpy()
map4 = cube2[:, :, 3].to_sunpy()
mapcube_seq = NDCubeSequence([map1, map2, map3, map4], common_axis=0)
@pytest.mark.parametrize("test_input,expected", [
def generate_ndcube(obs="example.fits.gz",
pf=None,
bp=None,
input_wcs=None,
zscale=1,
nz=30,
finite_energy=True,
xrange=None,
yrange=None):
"""
Generate NDCubes from the potential field extrapolation generated by
MONAMI
Parameters
----------
obs : `string`
file name of the fits file of the observation, in the case of
non-linear force free field extrapolation, this file must be the
Br component of the observation
pf : `string`
file name of the field extrapolation result
bp : `dictionary`
output of function read_data(obs, pf)
input_wcs : `astropy.wcs.WCS`
output of funciton construct_wcs(bp)
zscale: `int`
Sets the z-scale (1 = same scale as the x,y axes before the
heliographic transformation). Default= 1
nz: `int`
Number of equally spaced grid points in the z direction
finite_energy: `bool`
whether or not to constrain with finite energy
xrange, yrange: `list`
with 2 elements, defines the edges of the ROI where the potential field
extrapolation will be performed
Returns
-------
`tuple`
with the first, second and third element as the NDCube for bxp, byp
and bzp
"""
# read data
if bp is None:
if pf is not None:
bp = read_data(obs, pf)
else:
# read fits file and get the fits header
hdul = fits.open(obs)
header = hdul[0].header
bz = hdul[0].data
hdul.close()
shape = np.shape(bz)
# get xrange and yrange
if xrange is None or yrange is None:
xrange, yrange = get_roi(bz)
# test if xrange and yrange is correct
if np.shape(xrange) != (2, ):
raise ValueError("xrange must be in form of [x0, x1]!")
if np.shape(yrange) != (2, ):
raise ValueError("yrange must be in form of [x0, x1]!")
# convert xrange and yrange into integers
xrange = np.rint(xrange).tolist()
yrange = np.rint(yrange).tolist()
xrange = list(map(int, xrange))
yrange = list(map(int, yrange))
# not beyond the boudnary
if xrange[1] >= shape[1]:
xrange[1] = shape[1] - 1
if yrange[1] >= shape[0]:
yrange[1] = shape[0] - 1
if xrange[0] < 0:
xrange[0] = 0
if yrange[0] < 0:
yrange[0] = 0
# make pixel numbers even
if (yrange[1] - yrange[0]) % 2 == 1:
yrange[1] = yrange[1] - 1
if (xrange[1] - xrange[0]) % 2 == 1:
xrange[1] = xrange[1] - 1
# call the potential field extrapolation code
print('Start to perform the potential field extrapolation...')
pfe = Magnetic_Field_Extrapolation(bz[yrange[0]:yrange[1],
xrange[0]:xrange[1]],
nz=nz,
zscale=zscale,
finite_energy=finite_energy)
# construct bp dictionary
bp = {
"bpx": pfe['Bx'],
"bpy": pfe['By'],
"bpz": pfe['Bz'],
"zscale": zscale,
"xrange": xrange,
"yrange": yrange,
"header": header
}
# construct wcs coordinate
if input_wcs is None:
input_wcs = construct_wcs(bp)
# create NDCubes
if 'bpx' in bp.keys():
metax = {"Description": "Bx magnetic field from MONAMI/PF"}
metay = {"Description": "By magnetic field from MONAMI/PF"}
metaz = {"Description": "Bz magnetic field from MONAMI/PF"}
bxtube = NDCube(bp['bpx'], input_wcs, meta=metax)
bytube = NDCube(bp['bpy'], input_wcs, meta=metay)
bztube = NDCube(bp['bpz'], input_wcs, meta=metaz)
else:
metax = {"Description": "Bx magnetic field from MONAMI/NLFFF"}
metay = {"Description": "By magnetic field from MONAMI/NLFFF"}
metaz = {"Description": "Bz magnetic field from MONAMI/NLFFF"}
bxtube = NDCube(bp['bnx'], input_wcs, meta=metax)
bytube = NDCube(bp['bny'], input_wcs, meta=metay)
bztube = NDCube(bp['bnz'], input_wcs, meta=metaz)
return (bxtube, bytube, bztube)