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)