Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Changed the axis order of coordsys conversion matrix / added doctring to image deconvolution classes #124

Merged
merged 12 commits into from
Jan 31, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
79 changes: 68 additions & 11 deletions cosipy/image_deconvolution/RichardsonLucy.py
Original file line number Diff line number Diff line change
@@ -1,14 +1,16 @@
import copy
import numpy as np
import astropy.units as u
from tqdm.autonotebook import tqdm
import gc

from histpy import Histogram

from .deconvolution_algorithm_base import DeconvolutionAlgorithmBase

class RichardsonLucy(DeconvolutionAlgorithmBase):
"""
A class for the RichardsonLucy algorithm.
The algorithm here is based on Knoedlseder+99, Knoedlseder+05, Siegert+20.
"""

def __init__(self, initial_model_map, data, parameter):
DeconvolutionAlgorithmBase.__init__(self, initial_model_map, data, parameter)
Expand Down Expand Up @@ -43,37 +45,61 @@ def __init__(self, initial_model_map, data, parameter):

self.smoothing_sigma = parameter['smoothing_FWHM'] / 2.354820 # degree

self.smoothing_max_sigma = parameter.get('smoothing_max_sigma', default = 5.0)
self.gaussian_filter = self.calc_gaussian_filter(self.smoothing_sigma, self.smoothing_max_sigma)
self.gaussian_filter = self.calc_gaussian_filter(self.smoothing_sigma)

def pre_processing(self):
pass

def Estep(self):
# self.expectation = self.calc_expectation(self.model_map, self.data, self.use_sparse)
"""
Notes
-----
Expect count histogram is calculated in the post processing.
"""
print("... skip E-step ...")

def Mstep(self):
"""
M-step in RL algorithm.

Notes
-----
Background normalization is also optimized based on a generalized RL algirithm.
Currenly we use a signle normalization parameter.
In the future, the normalization will be optimized for each background group defined in some file.
"""
# Currenly (2024-01-12) this method can work for both local coordinate CDS and in galactic coordinate CDS.
# This is just because in DC2 the rotate response for galactic coordinate CDS does not have an axis for time/scatt binning.
# However it is likely that it will have such an axis in the future in order to consider background variability depending on time and pointign direction etc.
# Then, the implementation here will not work. Thus, keep in mind that we need to modify it once the response format is fixed.

diff = self.data.event_dense / self.expectation - 1

delta_map_part1 = self.model_map / self.data.image_response_dense_projected
delta_map_part2 = Histogram(self.model_map.axes, unit = self.data.image_response_dense_projected.unit)

if self.data.response_on_memory == True:
diff_x_response_this_pix = np.tensordot(diff.contents, self.data.image_response_dense.contents, axes = ([1,2,3], [2,3,4]))
diff_x_response = np.tensordot(diff.contents, self.data.image_response_dense.contents, axes = ([1,2,3], [2,3,4]))
# [Time/ScAtt, Em, Phi, PsiChi] x [NuLambda, Ei, Em, Phi, PsiChi] -> [Time/ScAtt, NuLambda, Ei]

delta_map_part2[:] = np.tensordot(self.data.coordsys_conv_matrix.contents, diff_x_response_this_pix, axes = ([1,2], [0,1])) * diff_x_response_this_pix.unit * self.data.coordsys_conv_matrix.unit #lb, Ei
# [lb, Time/ScAtt, NuLambda] x [Time/ScAtt, NuLambda, Ei] -> [lb, Ei]
delta_map_part2[:] = np.tensordot(self.data.coordsys_conv_matrix.contents, diff_x_response, axes = ([0,2], [0,1])) \
* diff_x_response.unit * self.data.coordsys_conv_matrix.unit
# [Time/ScAtt, lb, NuLambda] x [Time/ScAtt, NuLambda, Ei] -> [lb, Ei]
# note that coordsys_conv_matrix is the sparse, so the unit should be recovered.

else:
for ipix in tqdm(range(self.npix_local)):
response_this_pix = np.sum(self.data.full_detector_response[ipix].to_dense(), axis = (4,5)) # may not work with the DC2 response format
if self.data.is_miniDC2_format == True:
response_this_pix = np.sum(self.data.full_detector_response[ipix].to_dense(), axis = (4,5)) # [Ei, Em, Phi, PsiChi]
else:
response_this_pix = self.data.full_detector_response[ipix].to_dense() # [Ei, Em, Phi, PsiChi]

diff_x_response_this_pix = np.tensordot(diff.contents, response_this_pix, axes = ([1,2,3], [1,2,3])) # Ti, Ei
diff_x_response_this_pix = np.tensordot(diff.contents, response_this_pix, axes = ([1,2,3], [1,2,3]))
# [Time/ScAtt, Em, Phi, PsiChi] x [Ei, Em, Phi, PsiChi] -> [Time/ScAtt, Ei]

delta_map_part2 += np.tensordot(self.data.coordsys_conv_matrix[:,:,ipix], diff_x_response_this_pix, axes = ([1],[0])) * diff_x_response_this_pix.unit * self.data.coordsys_conv_matrix.unit #lb, Ei
delta_map_part2 += np.tensordot(self.data.coordsys_conv_matrix[:,:,ipix], diff_x_response_this_pix, axes = ([0],[0])) \
* diff_x_response_this_pix.unit * self.data.coordsys_conv_matrix.unit #lb, Ei
# [Time/ScAtt, lb] x [Time/ScAtt, Ei] -> [lb, Ei]

self.delta_map = delta_map_part1 * delta_map_part2

Expand All @@ -86,6 +112,12 @@ def Mstep(self):
self.bkg_norm = self.bkg_norm_range[1]

def post_processing(self):
"""
Here three processes will be performed.
- response weighting filter: the delta map is renormalized as pixels with large exposure times will have more feedback.
- gaussian smoothing filter: the delta map is blurred with a Gaussian function.
- acceleration of RL algirithm: the normalization of delta map is increased as long as the updated image has no non-negative components.
"""

if self.do_response_weighting:
self.delta_map[:,:] *= self.response_weighting_filter
Expand All @@ -103,11 +135,28 @@ def post_processing(self):
self.expectation = self.calc_expectation(self.model_map, self.data)

def check_stopping_criteria(self, i_iteration):
"""
If i_iteration is smaller than iteration_max, the iterative process will continue.

Returns
-------
bool
"""
if i_iteration < self.iteration_max:
return False
return True

def register_result(self, i_iteration):
"""
The values below are stored at the end of each iteration.
- iteration: iteration number
- model_map: updated image
- delta_map: delta map after M-step
- processed_delta_map: delta map after post-processing
- alpha: acceleration parameter in RL algirithm
- background_normalization: optimized background normalization
- loglikelihood: log-likelihood
"""
loglikelihood = self.calc_loglikelihood(self.data, self.model_map, self.expectation)

this_result = {"iteration": i_iteration,
Expand Down Expand Up @@ -136,6 +185,14 @@ def show_result(self, i_iteration):
print(f' background_normalization: {self.result["background_normalization"]}')

def calc_alpha(self, delta, model_map, almost_zero = 1e-5): #almost_zero is needed to prevent producing a flux below zero
"""
Calculate the acceleration parameter in RL algorithm.

Returns
-------
float
Acceleration parameter
"""
alpha = -1.0 / np.min( delta / model_map ) * (1 - almost_zero)
alpha = min(alpha, self.alpha_max)
if alpha < 1.0:
Expand Down
154 changes: 97 additions & 57 deletions cosipy/image_deconvolution/coordsys_conversion_matrix.py
Original file line number Diff line number Diff line change
@@ -1,16 +1,18 @@
import astropy.units as u
from astropy.time import Time
from astropy.coordinates import SkyCoord, cartesian_to_spherical, Galactic
import numpy as np
import healpy as hp
from tqdm.autonotebook import tqdm
import sparse
import astropy.units as u
from astropy.time import Time
from astropy.coordinates import SkyCoord, cartesian_to_spherical, Galactic

from scoords import Attitude, SpacecraftFrame
from histpy import Histogram, Axes, Axis, HealpixAxis

import sparse

class CoordsysConversionMatrix(Histogram):
"""
A class for coordinate conversion matrix (ccm).
"""

def __init__(self, edges, contents = None, sumw2 = None,
labels=None, axis_scale = None, sparse = None, unit = None,
Expand All @@ -24,40 +26,50 @@ def __init__(self, edges, contents = None, sumw2 = None,
@classmethod
def time_binning_ccm(cls, full_detector_response, orientation, time_intervals, nside_model = None, is_nest_model = False):
"""
Calculate a ccm from a given orientation.

Parameters
----------
full_detector_response:
orientation:
time_intervals: 2d np.array. it is the same format of binned_data.axes['Time'].edges
nside_model: If it is None, it will be the same as the NSIDE in the response.
full_detector_response : :py:class:`cosipy.response.FullDetectorResponse`
Response
orientation : :py:class:`cosipy.spacecraftfile.SpacecraftFile`
Orientation
time_intervals : :py:class:`np.array`
The same format of binned_data.axes['Time'].edges
nside_model : int or None, default None
If it is None, it will be the same as the NSIDE in the response.
is_nest_model : bool, default False
If scheme of the model map is nested, it should be False while it is rare.

Returns
-------
coordsys_conv_matrix: Axes [ "lb", "Time", "NuLambda" ]
:py:class:`cosipy.image_deconvolution.CoordsysConversionMatrix`
Its axes are [ "Time", "lb", "NuLambda" ].
"""

if nside_model is None:
nside_model = full_detector_response.nside

axis_model_map = HealpixAxis(nside = nside_model, coordsys = "galactic", label = "lb")
axis_time = Axis(edges = time_intervals, label = "Time")
axis_model_map = HealpixAxis(nside = nside_model, coordsys = "galactic", label = "lb")
axis_local_map = full_detector_response.axes["NuLambda"]

axis_coordsys_conv_matrix = [ axis_model_map, axis_time, axis_local_map ] #lb, Time, NuLambda
axis_coordsys_conv_matrix = [ axis_time, axis_model_map, axis_local_map ] #Time, lb, NuLambda

contents = []

for ipix in tqdm(range(hp.nside2npix(nside_model))):
l, b = hp.pix2ang(nside_model, ipix, nest=is_nest_model, lonlat=True)
pixel_coord = SkyCoord(l, b, unit = "deg", frame = 'galactic')

ccm_thispix = np.zeros((axis_time.nbins, axis_local_map.nbins)) # without unit
for i_time, [init_time, end_time] in tqdm(enumerate(axis_time.bounds), total = len(axis_time.bounds)):
ccm_thispix = np.zeros((axis_model_map.nbins, axis_local_map.nbins)) # without unit

for i_time, [init_time, end_time] in enumerate(axis_time.bounds):
init_time = Time(init_time, format = 'unix')
end_time = Time(end_time, format = 'unix')
init_time = Time(init_time, format = 'unix')
end_time = Time(end_time, format = 'unix')

filtered_orientation = orientation.source_interval(init_time, end_time)
filtered_orientation = orientation.source_interval(init_time, end_time)

for ipix in range(hp.nside2npix(nside_model)):
l, b = hp.pix2ang(nside_model, ipix, nest=is_nest_model, lonlat=True)
pixel_coord = SkyCoord(l, b, unit = "deg", frame = 'galactic')

pixel_movement = filtered_orientation.get_target_in_sc_frame(target_name = f"pixel_{ipix}_{i_time}",
target_coord = pixel_coord,
quiet = True,
Expand All @@ -70,10 +82,10 @@ def time_binning_ccm(cls, full_detector_response, orientation, time_intervals, n
src_path = pixel_movement,
save = False)

ccm_thispix[i_time] = dwell_time_map.data
ccm_thispix[ipix] = dwell_time_map.data
# (HealpixMap).data returns the numpy array without its unit. dwell_time_map.unit is u.s.

ccm_thispix_sparse = sparse.COO.from_numpy( ccm_thispix.reshape((1, axis_time.nbins, axis_local_map.nbins)) )
ccm_thispix_sparse = sparse.COO.from_numpy( ccm_thispix.reshape((1, axis_model_map.nbins, axis_local_map.nbins)) )

contents.append(ccm_thispix_sparse)

Expand All @@ -86,15 +98,28 @@ def time_binning_ccm(cls, full_detector_response, orientation, time_intervals, n
@classmethod
def spacecraft_attitude_binning_ccm(cls, full_detector_response, exposure_table, nside_model = None, use_averaged_pointing = False):
"""
Calculate a ccm from a given exposure_table.

Parameters
----------
full_detector_response:
exposure_table:
use_averaged_pointing: if this is True, the ccm loses accuracy but the calculatiion gets much faster.
full_detector_response : :py:class:`cosipy.response.FullDetectorResponse`
Response
exposure_table : :py:class:`cosipy.image_deconvolution.SpacecraftAttitudeExposureTable`
Scatt exposure table
nside_model : int or None, default None
If it is None, it will be the same as the NSIDE in the response.
use_averaged_pointing : bool, default False
If it is True, first the averaged Z- and X-pointings are calculated.
Then the dwell time map is calculated once for ach model pixel and each scatt_binning_index.
If it is False, the dwell time map is calculated for each attitude in zpointing and xpointing in the exposure table.
Then the calculated dwell time maps are summed up.
In the former case, the computation is fast but may lose the angular resolution.
In the latter case, the conversion matrix is more accurate but it takes a long time to calculate it.

Returns
-------
coordsys_conv_matrix: Axes [ "lb", "ScAtt", "NuLambda" ]
:py:class:`cosipy.image_deconvolution.CoordsysConversionMatrix'
Its axes are [ "ScAtt", "lb", "NuLambda" ].
"""

if nside_model is None:
Expand All @@ -104,41 +129,41 @@ def spacecraft_attitude_binning_ccm(cls, full_detector_response, exposure_table,

n_scatt_bins = len(exposure_table)

axis_model_map = HealpixAxis(nside = nside_model, coordsys = "galactic", scheme = exposure_table.scheme, label = "lb")
axis_scatt = Axis(edges = np.arange(n_scatt_bins+1), label = "ScAtt")
axis_model_map = HealpixAxis(nside = nside_model, coordsys = "galactic", scheme = exposure_table.scheme, label = "lb")
axis_local_map = full_detector_response.axes["NuLambda"]

axis_coordsys_conv_matrix = [ axis_model_map, axis_scatt, axis_local_map ] #lb, ScAtt, NuLambda
axis_coordsys_conv_matrix = [ axis_scatt, axis_model_map, axis_local_map ] #lb, ScAtt, NuLambda

contents = []

for ipix in tqdm(range(hp.nside2npix(nside_model))):
l, b = hp.pix2ang(nside_model, ipix, nest=is_nest_model, lonlat=True)
pixel_coord = SkyCoord(l, b, unit = "deg", frame = 'galactic')

ccm_thispix = np.zeros((axis_scatt.nbins, axis_local_map.nbins)) # without unit
for i_scatt_bin in tqdm(range(n_scatt_bins)):
ccm_thispix = np.zeros((axis_model_map.nbins, axis_local_map.nbins)) # without unit

for idx in range(n_scatt_bins):
row = exposure_table.iloc[idx]

scatt_binning_index = row['scatt_binning_index']
num_pointings = row['num_pointings']
#healpix_index = row['healpix_index']
zpointing = row['zpointing']
xpointing = row['xpointing']
zpointing_averaged = row['zpointing_averaged']
xpointing_averaged = row['xpointing_averaged']
delta_time = row['delta_time']
exposure = row['exposure']

if use_averaged_pointing:
z = SkyCoord([zpointing_averaged[0]], [zpointing_averaged[1]], frame="galactic", unit="deg")
x = SkyCoord([xpointing_averaged[0]], [xpointing_averaged[1]], frame="galactic", unit="deg")
else:
z = SkyCoord(zpointing.T[0], zpointing.T[1], frame="galactic", unit="deg")
x = SkyCoord(xpointing.T[0], xpointing.T[1], frame="galactic", unit="deg")
row = exposure_table.iloc[i_scatt_bin]

scatt_binning_index = row['scatt_binning_index']
num_pointings = row['num_pointings']
#healpix_index = row['healpix_index']
zpointing = row['zpointing']
xpointing = row['xpointing']
zpointing_averaged = row['zpointing_averaged']
xpointing_averaged = row['xpointing_averaged']
delta_time = row['delta_time']
exposure = row['exposure']

attitude = Attitude.from_axes(x = x, z = z, frame = 'galactic')
if use_averaged_pointing:
z = SkyCoord([zpointing_averaged[0]], [zpointing_averaged[1]], frame="galactic", unit="deg")
x = SkyCoord([xpointing_averaged[0]], [xpointing_averaged[1]], frame="galactic", unit="deg")
else:
z = SkyCoord(zpointing.T[0], zpointing.T[1], frame="galactic", unit="deg")
x = SkyCoord(xpointing.T[0], xpointing.T[1], frame="galactic", unit="deg")

attitude = Attitude.from_axes(x = x, z = z, frame = 'galactic')

for ipix in range(hp.nside2npix(nside_model)):
l, b = hp.pix2ang(nside_model, ipix, nest=is_nest_model, lonlat=True)
pixel_coord = SkyCoord(l, b, unit = "deg", frame = 'galactic')

src_path_cartesian = SkyCoord(np.dot(attitude.rot.inv().as_matrix(), pixel_coord.cartesian.xyz.value),
representation_type = 'cartesian', frame = SpacecraftFrame())
Expand All @@ -159,9 +184,9 @@ def spacecraft_attitude_binning_ccm(cls, full_detector_response, exposure_table,

hist, bins = np.histogram(pixels, bins = axis_local_map.edges, weights = weights)

ccm_thispix[idx] = hist
ccm_thispix[ipix] = hist

ccm_thispix_sparse = sparse.COO.from_numpy( ccm_thispix.reshape((1, axis_scatt.nbins, axis_local_map.nbins)) )
ccm_thispix_sparse = sparse.COO.from_numpy( ccm_thispix.reshape((1, axis_model_map.nbins, axis_local_map.nbins)) )

contents.append(ccm_thispix_sparse)

Expand All @@ -173,12 +198,27 @@ def spacecraft_attitude_binning_ccm(cls, full_detector_response, exposure_table,

@classmethod
def open(cls, filename, name = 'hist'):
"""
Open a ccm from a file.

Parameters
----------
filename : str
Path to file.
name : str, default 'hist'
Name of group where the histogram was saved.

Returns
-------
:py:class:`cosipy.image_deconvolution.CoordsysConversionMatrix'
Its axes are [ "lb", "Time" or "ScAtt", "NuLambda" ].
"""

new = super().open(filename, name)

new = cls(new.axes, contents = new.contents, sumw2 = new.contents, unit = new.unit)

new.binning_method = new.axes.labels[1] # 'Time' or 'ScAtt'
new.binning_method = new.axes.labels[0] # 'Time' or 'ScAtt'

return new

Expand Down
Loading
Loading