class DispersionDMX(Dispersion):
"""This class provides a DMX model - multiple DM values.
This model lets the user specify time ranges and fit for a different
DM value in each time range.
"""
register = True
category = "dispersion_dmx"
def __init__(self):
super(DispersionDMX, self).__init__()
# DMX is for info output right now
self.add_param(
floatParameter(
name="DMX",
units="pc cm^-3",
value=0.0,
description="Dispersion measure",
))
self.add_param(
prefixParameter(
name="DMX_0001",
units="pc cm^-3",
value=0.0,
unit_template=lambda x: "pc cm^-3",
description="Dispersion measure variation",
description_template=lambda x: "Dispersion measure",
paramter_type="float",
))
self.add_param(
prefixParameter(
name="DMXR1_0001",
units="MJD",
unit_template=lambda x: "MJD",
description="Beginning of DMX interval",
description_template=lambda x: "Beginning of DMX interval",
parameter_type="MJD",
time_scale="utc",
))
self.add_param(
prefixParameter(
name="DMXR2_0001",
units="MJD",
unit_template=lambda x: "MJD",
description="End of DMX interval",
description_template=lambda x: "End of DMX interval",
parameter_type="MJD",
time_scale="utc",
))
self.dm_value_funcs += [self.dmx_dm]
self.set_special_params(["DMX_0001", "DMXR1_0001", "DMXR2_0001"])
self.delay_funcs_component += [self.DMX_dispersion_delay]
def setup(self):
super(DispersionDMX, self).setup()
# Get DMX mapping.
# Register the DMX derivatives
for prefix_par in self.get_params_of_type("prefixParameter"):
if prefix_par.startswith("DMX_"):
self.register_deriv_funcs(self.d_delay_d_dmparam, prefix_par)
self.register_dm_deriv_funcs(self.d_dm_d_DMX, prefix_par)
def validate(self):
""" Validate the DMX parameters.
"""
super(DispersionDMX, self).validate()
DMX_mapping = self.get_prefix_mapping_component("DMX_")
DMXR1_mapping = self.get_prefix_mapping_component("DMXR1_")
DMXR2_mapping = self.get_prefix_mapping_component("DMXR2_")
if len(DMX_mapping) != len(DMXR1_mapping):
errorMsg = "Number of DMX_ parameters is not"
errorMsg += "equals to Number of DMXR1_ parameters. "
errorMsg += "Please check your prefixed parameters."
raise AttributeError(errorMsg)
if len(DMX_mapping) != len(DMXR2_mapping):
errorMsg = "Number of DMX_ parameters is not"
errorMsg += "equals to Number of DMXR2_ parameters. "
errorMsg += "Please check your prefixed parameters."
raise AttributeError(errorMsg)
def dmx_dm(self, toas):
condition = {}
tbl = toas.table
if not hasattr(self, "dmx_toas_selector"):
self.dmx_toas_selector = TOASelect(is_range=True)
DMX_mapping = self.get_prefix_mapping_component("DMX_")
DMXR1_mapping = self.get_prefix_mapping_component("DMXR1_")
DMXR2_mapping = self.get_prefix_mapping_component("DMXR2_")
for epoch_ind in DMX_mapping.keys():
r1 = getattr(self, DMXR1_mapping[epoch_ind]).quantity
r2 = getattr(self, DMXR2_mapping[epoch_ind]).quantity
condition[DMX_mapping[epoch_ind]] = (r1.mjd, r2.mjd)
select_idx = self.dmx_toas_selector.get_select_index(
condition, tbl["mjd_float"])
# Get DMX delays
dm = np.zeros(len(tbl)) * self.DM.units
for k, v in select_idx.items():
dm[v] = getattr(self, k).quantity
return dm
def DMX_dispersion_delay(self, toas, acc_delay=None):
""" This is a wrapper function for interacting with the TimingModel class
"""
return self.dispersion_type_delay(toas)
def d_dm_d_DMX(self, toas, param_name, acc_delay=None):
condition = {}
tbl = toas.table
if not hasattr(self, "dmx_toas_selector"):
self.dmx_toas_selector = TOASelect(is_range=True)
param = getattr(self, param_name)
dmx_index = param.index
DMXR1_mapping = self.get_prefix_mapping_component("DMXR1_")
DMXR2_mapping = self.get_prefix_mapping_component("DMXR2_")
r1 = getattr(self, DMXR1_mapping[dmx_index]).quantity
r2 = getattr(self, DMXR2_mapping[dmx_index]).quantity
condition = {param_name: (r1.mjd, r2.mjd)}
select_idx = self.dmx_toas_selector.get_select_index(
condition, tbl["mjd_float"])
try:
bfreq = self.barycentric_radio_freq(toas)
except AttributeError:
warn("Using topocentric frequency for dedispersion!")
bfreq = tbl["freq"]
dmx = np.zeros(len(tbl))
for k, v in select_idx.items():
dmx[v] = 1.0
return dmx * (u.pc / u.cm**3) / (u.pc / u.cm**3)
def print_par(self, ):
result = ""
DMX_mapping = self.get_prefix_mapping_component("DMX_")
DMXR1_mapping = self.get_prefix_mapping_component("DMXR1_")
DMXR2_mapping = self.get_prefix_mapping_component("DMXR2_")
result += getattr(self, "DMX").as_parfile_line()
sorted_list = sorted(DMX_mapping.keys())
for ii in sorted_list:
result += getattr(self, DMX_mapping[ii]).as_parfile_line()
result += getattr(self, DMXR1_mapping[ii]).as_parfile_line()
result += getattr(self, DMXR2_mapping[ii]).as_parfile_line()
return result
class DispersionDMX(Dispersion):
"""This class provides a DMX model - multiple DM values.
This model lets the user specify time ranges and fit for a different
DM value in each time range.
Parameters supported:
.. paramtable::
:class: pint.models.dispersion_model.DispersionDMX
"""
register = True
category = "dispersion_dmx"
def __init__(self):
super().__init__()
# DMX is for info output right now
self.add_param(
floatParameter(
name="DMX",
units="pc cm^-3",
value=0.0,
description="Dispersion measure",
))
self.add_DMX_range(None, None, dmx=0, frozen=False, index=1)
self.dm_value_funcs += [self.dmx_dm]
self.set_special_params(["DMX_0001", "DMXR1_0001", "DMXR2_0001"])
self.delay_funcs_component += [self.DMX_dispersion_delay]
def add_DMX_range(self,
mjd_start,
mjd_end,
index=None,
dmx=0,
frozen=True):
"""Add DMX range to a dispersion model with specified start/end MJDs and DMX.
Parameters
----------
mjd_start : float
MJD for beginning of DMX event.
mjd_end : float
MJD for end of DMX event.
index : int, None
Integer label for DMX event. If None, will increment largest used index by 1.
dmx : float
Change in DM during DMX event.
frozen : bool
Indicates whether DMX will be fit.
Returns
-------
index : int
Index that has been assigned to new DMX event.
"""
#### Setting up the DMX title convention. If index is None, want to increment the current max DMX index by 1.
if index is None:
dct = self.get_prefix_mapping_component("DMX_")
index = np.max(list(dct.keys())) + 1
i = f"{int(index):04d}"
if mjd_end is not None and mjd_start is not None:
if mjd_end < mjd_start:
raise ValueError("Starting MJD is greater than ending MJD.")
elif mjd_start != mjd_end:
raise ValueError("Only one MJD bound is set.")
if int(index) in self.get_prefix_mapping_component("DMX_"):
raise ValueError(
"Index '%s' is already in use in this model. Please choose another."
% index)
self.add_param(
prefixParameter(
name="DMX_" + i,
units="pc cm^-3",
value=dmx,
description="Dispersion measure variation",
parameter_type="float",
frozen=frozen,
))
self.add_param(
prefixParameter(
name="DMXR1_" + i,
units="MJD",
description="Beginning of DMX interval",
parameter_type="MJD",
time_scale="utc",
value=mjd_start,
))
self.add_param(
prefixParameter(
name="DMXR2_" + i,
units="MJD",
description="End of DMX interval",
parameter_type="MJD",
time_scale="utc",
value=mjd_end,
))
self.setup()
self.validate()
return index
def remove_DMX_range(self, index):
"""Removes all DMX parameters associated with a given index/list of indices.
Parameters
----------
index : float, int, list, np.ndarray
Number or list/array of numbers corresponding to DMX indices to be removed from model.
"""
if (isinstance(index, int) or isinstance(index, float)
or isinstance(index, np.int64)):
indices = [index]
elif not isinstance(index, list) or not isinstance(index, np.ndarray):
raise TypeError(
f"index must be a float, int, list, or array - not {type(index)}"
)
for index in indices:
index_rf = f"{int(index):04d}"
for prefix in ["DMX_", "DMXR1_", "DMXR2_"]:
self.remove_param(prefix + index_rf)
self.validate()
def get_indices(self):
"""Returns an array of integers corresponding to DMX parameters.
Returns
-------
inds : np.ndarray
Array of DMX indices in model.
"""
inds = []
for p in self.params:
if "DMX_" in p:
inds.append(int(p.split("_")[-1]))
return np.array(inds)
def setup(self):
super().setup()
# Get DMX mapping.
# Register the DMX derivatives
for prefix_par in self.get_params_of_type("prefixParameter"):
if prefix_par.startswith("DMX_"):
self.register_deriv_funcs(self.d_delay_d_dmparam, prefix_par)
self.register_dm_deriv_funcs(self.d_dm_d_DMX, prefix_par)
def validate(self):
"""Validate the DMX parameters."""
super().validate()
DMX_mapping = self.get_prefix_mapping_component("DMX_")
DMXR1_mapping = self.get_prefix_mapping_component("DMXR1_")
DMXR2_mapping = self.get_prefix_mapping_component("DMXR2_")
if DMX_mapping.keys() != DMXR1_mapping.keys():
# FIXME: report mismatch
raise ValueError("DMX_ parameters do not "
"match DMXR1_ parameters. "
"Please check your prefixed parameters.")
if DMX_mapping.keys() != DMXR2_mapping.keys():
raise ValueError("DMX_ parameters do not "
"match DMXR2_ parameters. "
"Please check your prefixed parameters.")
def validate_toas(self, toas):
DMX_mapping = self.get_prefix_mapping_component("DMX_")
DMXR1_mapping = self.get_prefix_mapping_component("DMXR1_")
DMXR2_mapping = self.get_prefix_mapping_component("DMXR2_")
bad_parameters = []
for k in DMXR1_mapping.keys():
if self._parent[DMX_mapping[k]].frozen:
continue
b = self._parent[DMXR1_mapping[k]].quantity.mjd * u.d
e = self._parent[DMXR2_mapping[k]].quantity.mjd * u.d
mjds = toas.get_mjds()
n = np.sum((b <= mjds) & (mjds < e))
if n == 0:
bad_parameters.append(DMX_mapping[k])
if bad_parameters:
raise MissingTOAs(bad_parameters)
def dmx_dm(self, toas):
condition = {}
tbl = toas.table
if not hasattr(self, "dmx_toas_selector"):
self.dmx_toas_selector = TOASelect(is_range=True)
DMX_mapping = self.get_prefix_mapping_component("DMX_")
DMXR1_mapping = self.get_prefix_mapping_component("DMXR1_")
DMXR2_mapping = self.get_prefix_mapping_component("DMXR2_")
for epoch_ind in DMX_mapping.keys():
r1 = getattr(self, DMXR1_mapping[epoch_ind]).quantity
r2 = getattr(self, DMXR2_mapping[epoch_ind]).quantity
condition[DMX_mapping[epoch_ind]] = (r1.mjd, r2.mjd)
select_idx = self.dmx_toas_selector.get_select_index(
condition, tbl["mjd_float"])
# Get DMX delays
dm = np.zeros(len(tbl)) * self._parent.DM.units
for k, v in select_idx.items():
dm[v] = getattr(self, k).quantity
return dm
def DMX_dispersion_delay(self, toas, acc_delay=None):
"""This is a wrapper function for interacting with the TimingModel class"""
return self.dispersion_type_delay(toas)
def d_dm_d_DMX(self, toas, param_name, acc_delay=None):
condition = {}
tbl = toas.table
if not hasattr(self, "dmx_toas_selector"):
self.dmx_toas_selector = TOASelect(is_range=True)
param = getattr(self, param_name)
dmx_index = param.index
DMXR1_mapping = self.get_prefix_mapping_component("DMXR1_")
DMXR2_mapping = self.get_prefix_mapping_component("DMXR2_")
r1 = getattr(self, DMXR1_mapping[dmx_index]).quantity
r2 = getattr(self, DMXR2_mapping[dmx_index]).quantity
condition = {param_name: (r1.mjd, r2.mjd)}
select_idx = self.dmx_toas_selector.get_select_index(
condition, tbl["mjd_float"])
try:
bfreq = self._parent.barycentric_radio_freq(toas)
except AttributeError:
warn("Using topocentric frequency for dedispersion!")
bfreq = tbl["freq"]
dmx = np.zeros(len(tbl))
for k, v in select_idx.items():
dmx[v] = 1.0
return dmx * (u.pc / u.cm**3) / (u.pc / u.cm**3)
def print_par(self):
result = ""
DMX_mapping = self.get_prefix_mapping_component("DMX_")
DMXR1_mapping = self.get_prefix_mapping_component("DMXR1_")
DMXR2_mapping = self.get_prefix_mapping_component("DMXR2_")
result += getattr(self, "DMX").as_parfile_line()
sorted_list = sorted(DMX_mapping.keys())
for ii in sorted_list:
result += getattr(self, DMX_mapping[ii]).as_parfile_line()
result += getattr(self, DMXR1_mapping[ii]).as_parfile_line()
result += getattr(self, DMXR2_mapping[ii]).as_parfile_line()
return result