"""Class representing an astronomical image (either real or synthetic), with plotting functionality."""

def __init__(self, pfov, image, image_hires = None, hires_scale = np.nan):

self.pfov = pfov #pixel field of view in arcsec

self.image = image

self.image_hires = image_hires #optional high-resolution version of the image

self.hires_scale = hires_scale

def __sub__(self, other):

"""Return a basic Plottable object whose image is the difference of self and other.

Note that the resulting object will not have an associated high-resolution image."""

if not np.isclose(self.pfov, other.pfov):

raise Exception("Tried to subtract two Plottables with different pixel sizes "

f"({self.pfov:.2f} / {other.pfov:.2f})")

else:

return Plottable(self.pfov, self.image - other.image)

def _projected_sep_array(self, centre):

"""Make an array of projected separations from a particular point to each image pixel.

Parameters

----------

centre : 2D tuple of ints

Indices of the point from which to calculate separations.

Returns

-------

2D array

Array of projected separations.

"""

dx, dy = np.meshgrid(np.arange(self.image.shape[1]) - centre[1],

np.arange(self.image.shape[0]) - centre[0])

return self.pfov * np.sqrt(dx**2 + dy**2)

def consistent_gaussian(self, radius = np.nan):

"""Establish whether self.image appears consistent with Gaussian noise.

If a radius is provided, return True if either the whole image, or the region within

radius arcsec of the centre, is consistent with Gaussian noise. The idea here is that

checking only the central region reduces the likelihood of other sources influencing

the result, but reducing the number of pixels also increases the significance of any

bright pixels, so it's useful to check both regions."""

if not np.isnan(radius):

sky_separation = self._projected_sep_array([i/2 for i in self.image.shape])

data = self.image[sky_separation < radius]

else:

data = self.image.flatten()

#perform an Anderson-Darling test for normality

result = anderson(data)

#according to scipy documentation, [2] should correspond to the 5% level,

#however this function is set up to account for possible future changes

sig = result.significance_level[2]

crit = result.critical_values[2]

#return the significance level and whether the data are consistent with a Gaussian at that level

if np.isnan(radius):

return sig, (result.statistic < crit)

else:

return sig, (result.statistic < crit or self.consistent_gaussian(np.nan)[1])

def _find_brightest(self, sep_threshold, centre):

"""Find the brightest pixel within a specified angular distance of a given pixel in self.image.

Parameters

----------

sep_threshold : float

Radius of search region in arcsec.

centre : 2D tuple of ints

Indices of centre of search region.

Returns

-------

2D tuple of ints

Indices of the brightest pixel.

"""

sky_separation = self._projected_sep_array(centre)

return np.unravel_index(np.ma.MaskedArray(self.image, sky_separation > sep_threshold).argmax(),

self.image.shape)

def _get_limits(self):

"""Calculate the appropriate axis limits in arcsec for a plot of the object's image."""

return [self.image.shape[1] * self.pfov / 2, -self.image.shape[1] * self.pfov / 2,

-self.image.shape[0] * self.pfov / 2, self.image.shape[0] * self.pfov / 2]

def plot(self, ax, plot_hires = False, xlabel = True, ylabel = True, log = False,

annotation = '', cmap_name = 'inferno', scalebar = False, dist = np.nan):

"""Plot self.image or self.image_hires using the provided axes, and add a colorbar below."""

#change units from mJy/pix to mJy/arcsec^2

intensity_scale = ((self.hires_scale if plot_hires else 1)/ self.pfov)**2

image = self.image_hires if plot_hires else self.image

if xlabel: ax.set_xlabel('$\mathregular{RA\ offset\ /\ arcsec}$')

if ylabel: ax.set_ylabel('$\mathregular{Dec\ offset\ /\ arcsec}$')

ax.tick_params(direction = 'in', color = 'white', width = 1, right = True, top = True)

limits = self._get_limits()

#plot NaN pixels (which will likely arise if plotting on a log scale) in black

cmap = plt.cm.get_cmap(cmap_name)

cmap.set_bad(color = 'k')

im = ax.imshow(np.log10(image * intensity_scale) if log else image * intensity_scale,

origin = 'lower',

interpolation = 'none', cmap = cmap,

extent = limits)

#put an annotation at the top left corner

ax.annotate(annotation, xy = (0.05, 0.95), xycoords = 'axes fraction', color = 'white',

verticalalignment = 'top', horizontalalignment = 'left')

#add a scalebar if desired

if scalebar:

if np.isnan(dist):

warnings.warn("No distance provided to plot_image. Unable to plot a scale bar.",

stacklevel = 2)

else:

#place scalebar at the lower left corner

scalebar_x = 0.05 #axis fraction

scalebar_y = 0.05 #axis fraction

#choose an appropriate scalebar length

if dist < 5:

scalebar_au = 10

elif dist < 1000:

scalebar_au = 100

else:

scalebar_au = 1000

scalebar_arcsec = scalebar_au / dist

ax.plot([limits[0] + scalebar_x * (limits[1] - limits[0]),

limits[0] + scalebar_x * (limits[1] - limits[0]) - scalebar_arcsec],

[limits[2] + scalebar_y * (limits[3] - limits[2]) for i in range(2)],

color = 'white')

ax.annotate(f'{scalebar_au} au',

xy = (scalebar_x + scalebar_arcsec / abs(limits[1] - limits[0]) + 0.02,

scalebar_y),

xycoords = 'axes fraction', color = 'white',

verticalalignment = 'center', horizontalalignment = 'left')

#add a colorbar

if not log: cblabel = '$\mathregular{Intensity\ /\ (mJy\ arcsec^{-2})}$'

else: cblabel = '$\mathregular{log\ [\ Intensity\ /\ (mJy\ arcsec^{-2})\ ]}$'

divider = make_axes_locatable(ax)

cax = divider.append_axes('bottom', size = '5%', pad = 0.6)

cb = plt.colorbar(im, cax = cax, orientation = 'horizontal')

cb.set_label(cblabel)

cb.ax.xaxis.set_major_formatter(FuncFormatter(lambda x , pos: f'{x:g}'))

cb.ax.xaxis.set_minor_locator(plt.NullLocator())

cb.ax.xaxis.set_major_locator(MaxNLocator(nbins = 5))

# show SIMBAD sources

source_coords = getattr(self, 'source_coords', [])

source_names = getattr(self, 'source_names', [])

for i, c in enumerate(source_coords):

#need a half-pixel offset to align correctly with the image

coord_x = limits[0] - self.pfov * (c[1] + 0.5)

coord_y = limits[2] + self.pfov * (c[0] + 0.5)

ax.plot([coord_x], [coord_y], c = 'tab:gray', marker = '+', markersize = 6)

ax.annotate(source_names[i], xy = (coord_x, coord_y),

xytext = (2, 2), textcoords = 'offset points',

color = 'tab:gray')

#return the relevant AxesImage

return im

def plot_contours(self, ax, rms, levels = [-3, -2, 2, 3], neg_col = 'gainsboro', pos_col = 'k'):

"""Plot contours showing the specified RMS levels of self.image."""

limits = self._get_limits()

#return the relevant QuadContourSet

return ax.contour(np.linspace(limits[0], limits[1], self.image.shape[1]),

np.linspace(limits[2], limits[3], self.image.shape[0]),

self.image, [i * rms for i in levels],

colors = [neg_col, neg_col, pos_col, pos_col], linestyles = 'solid')

def shifted(self, params):

"""Return a shifted and scaled copy of self.image."""

"""Enumeration used to specify the method used for making a model image."""

Geometric = 0

"""Class representing a disc model, with capability to make synthetic images of the modelled disc."""

def __init__(self, params, shape, pfov, aupp, hires_scale, alpha, include_unres,

stellarflux, flux_factor, model_type = ModelType.Geometric, npart = 100000):

"""Store the defining properties of the model."""

if model_type == ModelType.Geometric:

self._flux_function = self._geometric_model

elif model_type == ModelType.Particle:

self._flux_function = self._particle_model

self.npart = npart

else:

raise Exception(f"Invalid model type: {model_type}")

(self.funres, self.fres, self.x0, self.y0,

self.r1, self.r2, self.inc, self.theta) = params if include_unres else np.concatenate(([0], params))

self.pfov = pfov

self.aupp = aupp

self.hires_scale = hires_scale

self.alpha = alpha

self.stellarflux = stellarflux

self.flux_factor = flux_factor

self.shape = shape

self.model_type = model_type

def _geometric_model(self):

"""Make a grid of model fluxes using the geometric method."""

#note the use of +self.x0 but -self.y0, since RA increases to the left

dx, dy = np.meshgrid(np.linspace(-self.aupp * self.shape[1] / 2, self.aupp * self.shape[1] / 2,

num = self.shape[1] * self.hires_scale) + self.x0,

np.linspace(-self.aupp * self.shape[0] / 2, self.aupp * self.shape[0] / 2,

num = self.shape[0] * self.hires_scale) - self.y0)

#'primed' coordinates, i.e. in a frame rotated by theta

dxpr = dx * np.cos(np.deg2rad(self.theta)) + dy * np.sin(np.deg2rad(self.theta))

dypr = -dx * np.sin(np.deg2rad(self.theta)) + dy * np.cos(np.deg2rad(self.theta))

#distances from the star to each model pixel

r = np.sqrt(dypr**2 + (dxpr / np.cos(np.deg2rad(self.inc)))**2)

#set disc flux based on r^-alpha profile

in_disc = (r > self.r1) & (r < self.r2)

flux = np.zeros(r.shape)

flux[in_disc] = r[in_disc] ** -self.alpha

#ensure we don't divide by zero (if there's no flux, we don't need to normalize anyway)

#NOTE: normalization for the geometric model works differently to the particle model, in that

#the flux within the image field of view, rather than the total flux in the disc, is set to fres.

#these are equivalent as long as the disc lies entirely within the image, but not if it

#extends partially outside the bounds.

if np.sum(flux) > 0:

flux = self.fres * flux / np.sum(flux)

return flux

def _particle_model(self):

"""Make a grid of model fluxes using the particle-based method."""

#first get a variable uniformly distributed from 0-1

u = np.random.uniform(size = self.npart)

#transform from uniform distribution to linear distribution, so that the surface

#number density of particles is constant across the disc (since it's n(r)/2*pi*r*dr)

#see e.g. https://mathworld.wolfram.com/RandomNumber.html

number_index = 1

r = ((self.r2 ** (1 + number_index) - self.r1 ** (1 + number_index)) * u

+ self.r1 ** (1 + number_index)) ** (1 / (1 + number_index))

#give the particles random azimuthal angles

phi = np.random.uniform(0, 2 * np.pi, size = self.npart)

#coordinates in the frame aligned with the major & minor axes

dypr = r * np.cos(phi)

dxpr = r * np.sin(phi) * np.cos(np.deg2rad(self.inc))

#coordinates in the frame aligned with the image axes

dx = dxpr * np.cos(np.deg2rad(self.theta)) - dypr * np.sin(np.deg2rad(self.theta)) - self.x0

dy = dxpr * np.sin(np.deg2rad(self.theta)) + dypr * np.cos(np.deg2rad(self.theta)) + self.y0

#bin the particles and weight them to give a r^-alpha profile that sums to self.fres

flux = np.histogram2d(dy, dx,

weights = (self.fres / np.sum(r ** -self.alpha)) * r ** -self.alpha,

bins = [np.linspace(-self.shape[i] * self.aupp / 2,

self.shape[i] * self.aupp / 2,

self.shape[i] * self.hires_scale + 1) for i in range(2)])[0]

return flux

def _make_hires(self):

"""Store a high-res image without the central bright pixel in self.image_hires, and

return a grid of fluxes including that pixel, calculated using the appropriate method."""

flux = self._flux_function()

#scale down for lost flux

flux /= self.flux_factor

#at this point, store the flux without the central bright pixel in self.image_hires

#to improve contrast if plotted later

self.image_hires = flux.copy()

#central pixel gets additional flux from star plus any unresolved flux

flux[int((self.shape[0] * self.hires_scale / 2)

+ self.y0/(self.aupp/self.hires_scale)),

int((self.shape[1] * self.hires_scale / 2)

- self.x0/(self.aupp/self.hires_scale))] += (self.stellarflux + self.funres) / self.flux_factor

return flux

def make_images(self, psf):

"""Store appropriate images for analysis/plotting in self.image and self.image_hires.

Parameters

----------

psf : Plottable

PSF used to make synthetic observation. Must have an image_hires at the same scale as self.image_hires.

"""

if self.hires_scale != psf.hires_scale:

raise Exception(f"Model and PSF scales do not match ({self.hires_scale:.2f} / {psf.hires_scale:.2f})")

#convolve high-resolution model with high-resolution PSF; note that the call to

#_make_hires stores self.image_hires

convolved_hires = convolve_fft(self._make_hires(), psf.image_hires)

#rebin to lower-resolution image pixel size

"""Class representing a PACS observation. Stores the image and important associated metadata."""

def __init__(self, filename, search_radius = 5, target_ra = np.nan, target_dec = np.nan, dist = np.nan,

boxsize = 13, hires_scale = 1, rotate_to = np.nan, normalize = False, query_simbad = False):

"""Load in an image, store some important parameters and perform initial image processing."""

with fits.open(filename) as fitsfile:

self.image = fitsfile['image'].data * 1000 #image in mJy/pixel

self.pfov = fitsfile['image'].header['CDELT2'] * 3600 #pixel FOV in arcsec

self.wav = int(fitsfile['PRIMARY'].header['WAVELNTH']) #wavelength of observations

self.level = int(fitsfile['PRIMARY'].header['LEVEL']) #processing level (20 or 25)

self.name = fitsfile['PRIMARY'].header['OBJECT'] #target name

self.angle = fitsfile['PRIMARY'].header['POSANGLE'] #pointing position angle

#extract the obsid; the appropriate keyword seemingly depends on the processing level

try:

self.obsid = fitsfile['PRIMARY'].header['OBSID001'] #this works for level 2.5

except KeyError:

self.obsid = fitsfile['PRIMARY'].header['OBS_ID'] #and this for level 2

#get the expected star coordinates in pixels, if RA and dec were provided;

#otherwise, assume it's at the centre of the image

if np.isnan(target_ra) or np.isnan(target_dec):

star_expected = [i / 2 for i in self.image.shape]

else:

wcs = WCS(fitsfile['image'].header)

star_expected = np.flip(wcs.wcs_world2pix([[target_ra, target_dec]], 0)[0])

#extract coverage level, so that we can estimate the rms flux in a suitable region

cov = fitsfile['coverage'].data

#refuse to analyse 160 micron data (70/100 is always available and generally at higher S/N)

if self.wav != 70 and self.wav != 100:

raise Exception(f"Please provide a 70 or 100 μm image ({filename} is at {self.wav} μm)")

#factors to correct for flux lost during high-pass filtering (see Kennedy et al. 2012)

if self.wav == 70:

self.flux_factor = 1.16

elif self.wav == 100:

self.flux_factor = 1.19

#if no distance is supplied, simply set d = 1 pc so that separations will be in arcsec, not au;

#in_au can be stored in any saved output for future reference, and plots can be annotated with sep_unit,

#which is intended to be embedded in a LaTeX string

if np.isnan(dist):

distance_provided = False

dist = 1

self.sep_unit = r'^{\prime\prime}'

self.in_au = False

else:

distance_provided = True

self.sep_unit = r'\mathrm{au}'

self.in_au = True

#au per pixel at the distance of the target

self.aupp = self.pfov * dist

#clean up NaN pixels

self.image[np.isnan(self.image)] = 0

cov[np.isnan(cov)] = 0

#find the coordinates of the brightest pixel within search_radius arcsec

#of the specified RA and dec (or simply the centre)

brightest_pix = self._find_brightest(search_radius, star_expected)

#estimate the rms flux in a region defined by two conditions: coverage is above a

#specified level, and projected separation from the brightest pixel is above a certain level.

#NOTE: if the provided RA/dec are far from the image centre, the region defined by these

#conditions may not be the most appropriate (however, it's unlikely that we will be trying

#to fit a source near the edge of a map)

cov_threshold_rms = 0.6 #fraction of max coverage

sep_threshold_rms = 15 #arcsec

sky_separation = self._projected_sep_array(brightest_pix)

self.rms = self._estimate_background((cov > cov_threshold_rms * np.max(cov)) &

(sky_separation > sep_threshold_rms))

#need to scale up uncertainties since noise is correlated

natural_pixsize = 3.2 #always the case for PACS 70/100 micron images

self.uncert = self.rms * self._correlated_noise_factor(natural_pixsize)

if np.isnan(rotate_to):

#no rotation requested; simply crop down to the requested size

self._crop_image(brightest_pix, boxsize)

else:

#cut out a portion of the image with the brightest pixel at the centre - this step is necessary

#because after the rotation the brightest_pix coordinates will no longer be correct

self._crop_image(brightest_pix, 2 * boxsize)

#rotate to the requested position angle (necessary if using image as a PSF)

self.image = rotate(self.image, self.angle - rotate_to)

#now cut down to the requested size; note that we again look for the brightest pixel

#and put this in the centre, since the rotation may have introduced a small offset

self._crop_image(self._find_brightest(2 * self.pfov, [i / 2 for i in self.image.shape]), boxsize)

#normalize if requested

if normalize: self.image /= np.sum(self.image)

#rebin to a higher resolution if requested

self.hires_scale = hires_scale

if self.hires_scale >= 1:

self.image_hires = congrid(self.image,

[i * self.hires_scale for i in self.image.shape],

minusone = True)

#ensure that flux is conserved

self.image_hires *= np.sum(self.image) / np.sum(self.image_hires)

else:

raise Exception(f"hires_scale should be an integer >= 1")

if query_simbad:

if not np.isnan(rotate_to):

warnings.warn(f"SIMBAD source overplotting for rotated images is"

" not supported. Skipping query.", stacklevel = 2)

else:

self.source_coords=[]

self.source_names=[]

Simbad.add_votable_fields('pm','plx')

with fits.open(filename) as fitsfile:

qra = fitsfile['PRIMARY'].header['RA'] if np.isnan(target_ra) else target_ra

qdec = fitsfile['PRIMARY'].header['DEC'] if np.isnan(target_dec) else target_dec

coord = SkyCoord(ra = qra, dec = qdec, unit = (u.degree, u.degree), frame = 'icrs')

#find sources within a circle whose radius is half the cutout side length

r = Simbad.query_region(coordinates = coord, radius = boxsize * self.pfov * u.arcsec)

if len(r) > 0:

wcs = WCS(fitsfile['image'].header)

for i in range(len(r)):

#assume a distance of 50pc if none is available

qdist = 50 * u.pc

#preferentially use the supplied distance

if distance_provided: qdist = dist * u.pc

#otherwise, try to get one from SIMBAD

elif r[i]['PLX_VALUE'] > 0: distance = (1e3 / r[i]['PLX_VALUE']) * u.pc

if np.isfinite(r[i]['PMRA']) and np.isfinite(r[i]['PMDEC']):

#J2000 sky coordinates

s2000 = SkyCoord(r[i]['RA'].replace(' ',':')+' '+r[i]['DEC'].replace(' ',':'),

unit = (u.hour, u.degree), distance = qdist,

pm_ra_cosdec = r[i]['PMRA'] * u.mas / u.yr,

pm_dec = r[i]['PMDEC'] * u.mas / u.yr,

obstime = Time(2451545.0,format='jd'))

#apply proper motion correction to observation date

s = s2000.apply_space_motion(new_obstime = Time(fitsfile['PRIMARY'].header['DATE-OBS']))

else:

s = SkyCoord(r[i]['RA'].replace(' ',':')+' '+r[i]['DEC'].replace(' ',':'),

unit = (u.hour, u.degree))

#find the pixel corresponding to source i

coord = np.flip(wcs.wcs_world2pix([[s.ra.deg, s.dec.deg]], 0)[0])

#translate into image cutout coordinates

coord -= np.array(brightest_pix) - boxsize

#store the coordinates and source name as an attribute

append = True

#don't append duplicate coordinates (i.e. planets)

for c in self.source_coords:

if np.isclose(c, coord).all():

append = False

if append:

self.source_coords.append(coord)

self.source_names.append(r[i]['MAIN_ID'].decode())

def best_psf_subtraction(self, psf, param_limits):

"""Return the best-fitting PSF-subtracted image."""

if not np.isclose(self.pfov, psf.pfov):

raise Exception("best_psf_subtraction received a PSF with the wrong pixel size "

f"({self.pfov:.2f} / {psf.pfov:.2f})")

#note that shiftmax here is in PACS pixels

limits = [(-param_limits.shiftmax, param_limits.shiftmax), #x shift

(-param_limits.shiftmax, param_limits.shiftmax), #y shift

(0, 2 * np.amax(self.image))] #peak flux

result = differential_evolution(lambda p: np.sum(((self.image - psf.shifted(p)) / self.uncert) ** 2),

limits)

return Plottable(image = self.image - psf.shifted(result['x']), pfov = self.pfov)

def _crop_image(self, centre, boxscale):

"""Crop self.image such that the specified pixel is perfectly centred.

Parameters

----------

centre : 2D tuple of ints

Indices of pixel to place at the centre.

boxscale : int

Cut out a square of dimension (2 * boxscale + 1).

"""

#this function guarantees that the desired pixel will be exactly at the centre

#as long as boxscale < centre[i], which is useful for making centred plots

self.image = self.image[int(centre[0] - boxscale) : int(centre[0] + boxscale + 1),

int(centre[1] - boxscale) : int(centre[1] + boxscale + 1)]

def _estimate_background(self, condition = None, sigma_level = 3.0, tol = 1e-6, max_iter = 20):

"""Estimate the background RMS of self.image or self.image[condition] using an iterative method.

Parameters

----------

condition : 2D array, optional

Boolean array specifying which pixels of self.image to process. (default: None)

sigma_level : float

Number of standard deviations used to define outliers. (default: 3.0)

tol : float, optional

Fractional difference between RMS iterations at which to cease iterating. (default: 1e-6)

max_iter : int, optional

Maximum number of iterations. (default: 20)

Returns

-------

float

Estimated background RMS.

"""

data = self.image.flatten() if condition is None else self.image[condition]

rms = np.std(data)

#arbitrarily set error to >tol so that the iterations can get started

err = 2 * tol

i = 0

while (err > tol) and (i < max_iter):

rmsprev = rms

#discard data classified as outliers

data = data[data < sigma_level * rms]

rms = np.std(data)

err = abs((rms - rmsprev) / rmsprev)

i += 1

if (i == max_iter) and (err > tol):

warnings.warn(f"_estimate_background did not converge after {max_iter} iterations."

" You may wish to check the image for issues.", stacklevel = 2)

return rms

def _correlated_noise_factor(self, natural_pixsize):

"""Calculate the uncertainty scale factor for correlated noise from Fruchter & Hook (2002).

Parameters

----------

natural_pixsize : float

Natural pixel size of the image in arcsec.

Returns

-------

float

Uncertainty scale factor.

"""

r = natural_pixsize / self.pfov

if r >= 1.0:

return r / (1.0 - 1.0 / (3.0 * r))

else:

"""Subtract model from observations and calculate the chi-squared goodness of fit value."""

model = Model(params, obs.image.shape, obs.pfov, obs.aupp, psf.hires_scale,

alpha, include_unres, stellarflux, obs.flux_factor, model_type, npart)

#impose uniform priors within some ranges;

#note that the fluxes don't have an upper limit here, to allow for extremely bright cases

if (model.funres < 0 or model.fres < 0

or model.r1 <= 0 or model.r2 <= 0

or model.r1 >= model.r2

or model.r1 > param_limits.rmax or model.r2 > param_limits.rmax

or model.inc < 0 or model.inc > param_limits.imax

or abs(model.x0) > param_limits.shiftmax * obs.aupp

or abs(model.y0) > param_limits.shiftmax * obs.aupp

or abs(model.theta) > 90):

return np.inf

#force the disc to be at least a model pixel wide if using the geometric model, otherwise

#unphysical models with just a few pixels scattered around the image can result

if model_type == ModelType.Geometric:

dr_pix = (model.r2 - model.r1) * np.cos(np.deg2rad(model.inc)) * (psf.hires_scale / obs.aupp)

if (dr_pix <= 1):

return np.inf

model.make_images(psf)

lower_uncertainty = None, upper_uncertainty = None, model_consistent = None,

in_au = None, stellarflux = None):

"""Save the main results of the fit in a pickle file."""

dict = {

'resolved': resolved,

'include_unres': include_unres,

'max_likelihood': max_likelihood,

'median': median,

'lower_uncertainty': lower_uncertainty,

'upper_uncertainty': upper_uncertainty,

'model_consistent': model_consistent,

'in_au': in_au,

'stellarflux': stellarflux

}

with open(savepath + '/params.pickle', 'wb') as file:

pickle.dump(dict, file, protocol = pickle.HIGHEST_PROTOCOL)

stellarflux = 0, boxsize = 13, hires_scale = 3, alpha = 1.5, include_unres = False,

initial_steps = 100, nwalkers = 200, nsteps = 800, burn = 600, ra = np.nan,

dec = np.nan, test = False, model_type = ModelType.Particle, npart = 100000,

query_simbad = False):

"""Fit one image and save the output."""

#if given no stellar flux, force an unresolved component to be added

if (stellarflux == 0 or np.isnan(stellarflux)) and not include_unres:

include_unres = True

warnings.warn("No stellar flux was supplied. Forcing the model to include an unresolved flux.",

stacklevel = 2)

obs = Observation(name_image, target_ra = ra, target_dec = dec, dist = dist, boxsize = boxsize,

query_simbad = query_simbad)

#put the star name, distance, obsid/level & wavelength together into an annotation for the image

annotation = '\n'.join([f'{obs.wav} μm image (level {(obs.level / 10):g})',

f'ObsID: {obs.obsid}',

name if name != '' else obs.name])

if not np.isnan(dist):

annotation += f' (at {dist:.1f} pc)'

#if no PSF is provided, select one based on the processing level and wavelength

if name_psf == '':

if obs.wav != 70 and obs.wav != 100:

raise Exception(f'No {wav} μm PSF is provided by default')

if obs.level != 20 and obs.level != 25:

raise Exception(f'No level {level} PSF is provided by default')

default_psf = {

70: {20: 'psf/gamma_dra_70/1342217404/level2/HPPPMAPB/'

'hpacs1342217404_20hpppmapb_00_1469423089198.fits.gz',

25: 'psf/gamma_dra_70/1342217404/level2_5/HPPHPFMAPB/'

'hpacs_25HPPHPFMAPB_blue_1757_p5129_00_v1.0_1470980845846.fits.gz'},

100: {20: 'psf/gamma_dra_100/1342216069/level2/HPPPMAPB/'

'hpacs1342216069_20hpppmapb_00_1469417766626.fits.gz',

25: 'psf/gamma_dra_100/1342216069/level2_5/HPPHPFMAPB/'

'hpacs_25HPPHPFMAPB_green_1757_p5129_00_v1.0_1470967312171.fits.gz'}

}

name_psf = default_psf[obs.wav][obs.level]

#a boxsize of 13 should be large enough to cover the PSF - no need for a larger PSF

#even if the image is larger (this would slow down the convolution)

psf = Observation(name_psf, boxsize = 13, hires_scale = hires_scale, rotate_to = obs.angle,

normalize = True)

#however, we need to store a PSF with the same dimension as the image for the initial subtraction

psf_imagesize = Observation(name_psf, boxsize = boxsize, rotate_to = obs.angle, normalize = True)

#abort execution if the PSF pixel scale doesn't match that of the image

if not np.isclose(psf.pfov, obs.pfov):

raise Exception(f"PSF and image pixel sizes do not match ({psf.pfov:.2f} / {obs.pfov:.2f})")

#issue a warning if the image and PSF are at different wavelengths

if psf.wav != obs.wav:

warnings.warn("The wavelength of the supplied PSF does not match that of the image"

f" ({psf.wav} / {obs.wav})",

stacklevel = 2)

#ensure that the output directory exists

if not os.path.exists(savepath):

os.makedirs(savepath)

#upper limits on the model parameters

#NOTE: the radii are restricted to the half-diagonal length of the image - this is a quick

#way of ensuring that we don't end up with arbitrarily large discs that lie completely outside

#the image cutout

#NOTE: all limits are used for the initial differential evolution fit (since this method requires

#parameter ranges to be specified); for the MCMC fit, fmax is NOT imposed, but the other limits are

param_limits = ParamLimits(

fmax = 200000, #mJy

shiftmax = 5, #PACS pixels

rmax = min(obs.image.shape) * obs.aupp / np.sqrt(2), #au

imax = 90 if model_type == ModelType.Particle else 88 #deg

)

#if requested, first check whether the image is consistent with a PSF and skip the fit if possible

psfsub = obs.best_psf_subtraction(psf_imagesize, param_limits)

if test:

sig, is_noise = psfsub.consistent_gaussian(radius = 15)

if is_noise:

print(f"The PSF subtraction is consistent with Gaussian noise at the {sig:.0f}% level."

" There is likely not a resolved disc here. Skipping this system.")

print("Exporting image of PSF subtraction...")

#make a two-panel image: [data, psf subtraction]

fig, ax = plt.subplots(nrows = 1, ncols = 2, figsize = (9, 6), sharey = True)

#first plot the PACS data

obs.plot(ax[0], annotation = annotation)

#then the PSF subtraction

psfsub.plot(ax[1], ylabel = False, annotation = 'PSF subtraction')

psfsub.plot_contours(ax[1], obs.rms)

plt.tight_layout()

fig.savefig(savepath + '/image_model.png', dpi = 150)

plt.close(fig)

#save a pickle simply indicating that no disc was resolved

save_params(savepath, False)

return

else:

print(f"The PSF subtraction is not consistent with Gaussian noise at the {sig:.0f}% level."

" There may be a resolved disc here. Performing disc fit.")

#find best-fitting parameters using differential evolution, which searches for the

#global minimum within the parameter ranges specified by the arguments.

#format is [<funres,> fres, x0, y0, r1, r2, inc, theta]

search_space = [(0, param_limits.fmax), (0, param_limits.fmax),

(-param_limits.shiftmax * obs.aupp, param_limits.shiftmax * obs.aupp),

(-param_limits.shiftmax * obs.aupp, param_limits.shiftmax * obs.aupp),

(0, param_limits.rmax), (0, param_limits.rmax),

(0, param_limits.imax), (-90, 90)]

pnames = [r'$F_\mathrm{unres}\ /\ \mathrm{mJy}$', r'$F_\mathrm{res}\ /\ \mathrm{mJy}$',

fr'$x_0\ /\ {obs.sep_unit}$', fr'$y_0\ /\ {obs.sep_unit}$',

fr'$r_1\ /\ {obs.sep_unit}$', fr'$r_2\ /\ {obs.sep_unit}$',

r'$i\ /\ \mathrm{deg}$', r'$\theta\ /\ \mathrm{deg}$'] #parameter names for plot labels

#if not including an unresolved flux, remove the first element of the parameter list

if not include_unres:

search_space.pop(0)

pnames.pop(0)

pnames[0] = r'$F_\mathrm{disc}\ /\ \mathrm{mJy}$'

print("Finding a suitable initial model...")

pbar = tqdm.tqdm(total = initial_steps)

#set tol = 0 to ensure that DE runs for the prescribed number of steps & the progress bar works

res = differential_evolution(chi2, search_space,

args = (psf, alpha, include_unres, stellarflux,

obs, param_limits, model_type, npart),

updating = 'deferred', workers = -1, #use multiprocessing

tol = 0, popsize = 20, maxiter = initial_steps, polish = False,

callback = (lambda xk, convergence: pbar.update()))

pbar.close()

p0 = res['x']

ndim = p0.size

print("Running MCMC sampler...")

with Pool() as pool:

sampler = emcee.EnsembleSampler(nwalkers, ndim, log_probability,

args = (psf, alpha, include_unres, stellarflux,

obs, param_limits, model_type, npart),

pool = pool)

#initialize the walkers with an ndim-dimensional Gaussian distribution

pos = [p0 + p0 * 0.01 * np.random.randn(ndim) for i in range(nwalkers)]

#run MCMC

pos, prob, state = sampler.run_mcmc(pos, nsteps, progress = True)

print("Pickling samples...")

#extract and save the samples, excluding burn-in, for future use in e.g. corner plots

samples = sampler.get_chain(discard = burn, flat = True)

with open(savepath + '/samples.pickle','wb') as file:

pickle.dump(samples, file, protocol = pickle.HIGHEST_PROTOCOL)

print("Exporting plot of MCMC chains...")

#save a plot of the chains

fig, ax = plt.subplots(ndim + 1, figsize = (12, 16))

chain = sampler.get_chain()

#first plot the parameter chains

for i in range(ndim):

ax[i].plot(chain[:, :, i], c = 'k', alpha = 0.3)

ax[i].set_ylabel(pnames[i])

ax[i].xaxis.set_major_locator(plt.NullLocator())

#then the log-probability

ax[-1].plot(sampler.get_log_prob(), c = 'k', alpha = 0.3)

ax[-1].xaxis.set_major_locator(MaxNLocator(integer = True))

ax[-1].set_ylabel('log probability')

ax[-1].set_xlabel('Step number')

#formatting common to all subplots

for i in range(ndim + 1):

ax[i].axvspan(0, burn - 0.5, alpha = 0.1, color = 'k')

ax[i].set_xlim(0, nsteps - 1)

plt.tight_layout(h_pad = 0.5)

fig.savefig(savepath + '/chains.pdf')

fig.savefig(savepath + '/chains.png', dpi = 150)

plt.close(fig)

print("Exporting corner plot...")

#make the corner plot

fig = corner.corner(samples, quantiles = [0.16, 0.50, 0.84],

labels = pnames, show_titles = True, title_fmt = '.1f')

fig.savefig(savepath + '/corner.pdf')

fig.savefig(savepath + '/corner.png', dpi = 150)

plt.close(fig)

#get max-likelihood model parameters (no need to exclude burn-in for this, just want the best fit)

max_likelihood = sampler.flatchain[np.argmax(sampler.flatlnprobability), :]

#get median and 16th/84th percentile parameters

median = np.median(samples, axis = 0)

lower_uncertainty = median - np.percentile(samples, 16, axis = 0)

upper_uncertainty = np.percentile(samples, 84, axis = 0) - median

#now make a four-panel image: [data, psf subtraction, high-res max-likelihood model, residuals]

print("Exporting image of best-fit model...")

model = Model(max_likelihood, obs.image.shape, obs.pfov, obs.aupp, hires_scale, alpha,

include_unres, stellarflux, obs.flux_factor, model_type, npart)

model.make_images(psf)

fig, ax = plt.subplots(nrows = 1, ncols = 4, figsize = (18, 6), sharey = True)

#first plot the PACS data

obs.plot(ax[0], annotation = annotation)

#then the PSF subtraction

psfsub.plot(ax[1], ylabel = False, annotation = 'PSF subtraction')

psfsub.plot_contours(ax[1], obs.rms)

#now the high-res model

annotation_model = 'High-resolution model '

annotation_model += f'({npart} particles)' if model_type == ModelType.Particle else '(geometric)'

annotation_model += f'\nUnresolved component{" " if include_unres else " not "}included'

if not obs.in_au:

annotation_model += f'\nNo distance{" or stellar flux " if stellarflux == 0 else " "}provided'

elif stellarflux == 0:

annotation_model += '\nNo stellar flux provided'

model.plot(ax[2], ylabel = False, plot_hires = True, annotation = annotation_model,

scalebar = obs.in_au, dist = dist, log = True)

#finally, the model residuals

residual = obs - model