Skip to content

Commit

Permalink
TF doc and enabling SData subplot
Browse files Browse the repository at this point in the history
  • Loading branch information
@AlanLoh
AlanLoh committed Dec 13, 2023
1 parent e1c19ab commit 4a459fd
Show file tree
Hide file tree
Showing 3 changed files with 97 additions and 67 deletions.
2 changes: 1 addition & 1 deletion nenupy/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
__copyright__ = "Copyright 2023, nenupy"
__credits__ = ["Alan Loh"]
__license__ = "MIT"
__version__ = "2.6.0"
__version__ = "2.6.1"
__maintainer__ = "Alan Loh"
__email__ = "alan.loh@obspm.fr"

Expand Down
70 changes: 45 additions & 25 deletions nenupy/beamlet/sdata.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -546,11 +546,12 @@ def fbackground(self):

# --------------------------------------------------------- #
# ------------------------ Methods ------------------------ #
def plot(self, polarization=None, figname=None, db=True, **kwargs):
def plot(self, fig=None, ax=None, polarization=None, figname=None, db=True, **kwargs):
"""
kwargs keys: cmap, title, cblabel, figsize, altaza, vmin, vmax
"""
import matplotlib.pyplot as plt
import matplotlib.dates as mdates

if polarization is None:
pol_idx = 0
Expand All @@ -573,10 +574,17 @@ def plot(self, polarization=None, figname=None, db=True, **kwargs):
if 'figsize' not in kwargs.keys():
kwargs['figsize'] = (15, 10)

fig = plt.figure(figsize=kwargs['figsize'])
if ax is None:
# Create a full figure from the start
fig = plt.figure(figsize=kwargs['figsize'])
ax = fig.add_subplot()
else:
# Fill up the input ax with the plot
pass

if len(dynspec.shape) == 1:
if dynspec.size == self.datetime.size:
plt.plot(
ax.plot(
self.datetime,
dynspec
)
Expand All @@ -585,20 +593,20 @@ def plot(self, polarization=None, figname=None, db=True, **kwargs):
for ptime in ptimes:
if (ptime < self.datetime[0]) or (ptime > self.datetime[-1]):
continue
plt.axvline(ptime.datetime, linestyle='-.', color='black')
plt.ylim((kwargs.get('vmin', None), kwargs.get('vmax', None)))
plt.xlabel(
f'Time (since {self.time[0].isot})'
ax.axvline(ptime.datetime, linestyle='-.', color='black')
ax.ylim((kwargs.get('vmin', None), kwargs.get('vmax', None)))
ax.set_xlabel(
f'Time (UTC from {self.time[0].isot})'
)
plt.ylabel(kwargs['cblabel'])
ax.set_ylabel(kwargs['cblabel'])
elif dynspec.size == self.freq.size:
plt.plot(
ax.plot(
self.freq.to(u.MHz).value,
dynspec
)
plt.ylim((kwargs.get('vmin', None), kwargs.get('vmax', None)))
plt.xlabel('Frequency (MHz)')
plt.ylabel(kwargs['cblabel'])
ax.set_ylim((kwargs.get('vmin', None), kwargs.get('vmax', None)))
ax.set_xlabel('Frequency (MHz)')
ax.set_ylabel(kwargs['cblabel'])
else:
if 'vmin' not in kwargs.keys():
kwargs['vmin'] = np.nanpercentile(dynspec, 5)
Expand All @@ -612,7 +620,7 @@ def plot(self, polarization=None, figname=None, db=True, **kwargs):
kwargs['vmax'] = np.nanpercentile(dynspec, 95)
else:
pass
plt.pcolormesh(
im = ax.pcolormesh(
self.datetime,
self.freq.to(u.MHz).value,
dynspec,
Expand All @@ -626,16 +634,16 @@ def plot(self, polarization=None, figname=None, db=True, **kwargs):
for ptime in ptimes:
if (ptime < self.datetime[0]) or (ptime > self.datetime[-1]):
continue
plt.axvline(ptime.datetime, linestyle='-.', color='black')
cbar = plt.colorbar(pad=0.03)#format='%.1e')
ax.axvline(ptime.datetime, linestyle='-.', color='black')
cbar = plt.colorbar(im, pad=0.03)#format='%.1e')
cbar.set_label(kwargs['cblabel'])

if kwargs.get("overlay", None) is not None:
ax = plt.gca()
#ax = plt.gca()
xlim = ax.get_xlim()
ylim = ax.get_ylim()
overlay_time, overlay_frequency, overlay_values = kwargs["overlay"]
plt.pcolor(
ax.pcolor(
overlay_time.datetime,
overlay_frequency.to(u.MHz).value,
overlay_values,
Expand All @@ -655,17 +663,30 @@ def plot(self, polarization=None, figname=None, db=True, **kwargs):
ax.set_xlim(xlim)
ax.set_ylim(ylim)

plt.xlabel(
f'Time (since {self.time[0].isot})'
ax.set_xlabel(
f'Time (UTC from {self.time[0].isot})'
)
plt.ylabel('Frequency (MHz)')
plt.title(kwargs['title'])

ax.set_ylabel('Frequency (MHz)')
ax.set_title(kwargs['title'])

# Set x axis labels
# ax.xaxis.set_minor_locator(mdates.MinuteLocator(byminute=[15, 30, 45]))
# ax.xaxis.set_major_locator(mdates.HourLocator())
# hourFmt = mdates.DateFormatter("%H", usetex=True)
# ax.xaxis.set_major_formatter(hourFmt)
locator = ax.xaxis.set_major_locator(mdates.AutoDateLocator())
ax.xaxis.set_major_formatter(
mdates.ConciseDateFormatter(locator, show_offset=False)
)

# Save or show
if not (ax is None):
return

if figname is None:
plt.show()
elif figname.lower() == 'return':
return fig
elif (figname is None) or (figname == ""):
return
else:
fig.savefig(
figname,
Expand All @@ -674,7 +695,6 @@ def plot(self, polarization=None, figname=None, db=True, **kwargs):
bbox_inches='tight'
)
plt.close('all')
return


# --------------------------------------------------------- #
Expand Down
92 changes: 51 additions & 41 deletions nenupy/io/tf.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -113,17 +113,17 @@ def func_call(self) -> Callable:

@classmethod
def correct_bandpass(cls):
""":class:`~nenupy.io.tf.TFTask` calling :func:`~nenupy.io.tf_utils.correct_bandpass` to correct the polyphase-filter bandpass reponse.
"""
""":class:`~nenupy.io.tf.TFTask` calling :func:`~nenupy.io.tf_utils.correct_bandpass` to correct the polyphase-filter bandpass reponse."""

def wrapper_task(data, channels):
return utils.correct_bandpass(data=data, n_channels=channels)

return cls("Correct bandpass", wrapper_task, ["channels"])

@classmethod
def remove_channels(cls):
""":class:`~nenupy.io.tf.TFTask` calling :func:`~nenupy.io.tf_utils.remove_channels_per_subband` to set a list of sub-band channels to `NaN` values.
"""
""":class:`~nenupy.io.tf.TFTask` calling :func:`~nenupy.io.tf_utils.remove_channels_per_subband` to set a list of sub-band channels to `NaN` values."""

def wrapper_task(data, channels, remove_channels):
if (remove_channels is None) or (len(remove_channels) == 0):
return data
Expand Down Expand Up @@ -180,8 +180,8 @@ def wrapper_task(

@classmethod
def correct_faraday_rotation(cls):
""":class:`~nenupy.io.tf.TFTask` calling :func:`~nenupy.io.tf_utils.de_faraday_data` to correct for Faraday rotation for a given ``'rotation_measure'`` set in :attr:`~nenupy.io.tf.TFPipeline.parameters`.
"""
""":class:`~nenupy.io.tf.TFTask` calling :func:`~nenupy.io.tf_utils.de_faraday_data` to correct for Faraday rotation for a given ``'rotation_measure'`` set in :attr:`~nenupy.io.tf.TFPipeline.parameters`."""

def apply_faraday(frequency_hz, data, rotation_measure):
if rotation_measure is None:
return frequency_hz, data
Expand All @@ -197,17 +197,20 @@ def apply_faraday(frequency_hz, data, rotation_measure):
def de_disperse(cls):
""":class:`~nenupy.io.tf.TFTask` calling :func:`~nenupy.io.tf_utils.de_disperse_array` to de-disperse the data using the ``'dispersion_measure'`` set in :attr:`~nenupy.io.tf.TFPipeline.parameters`.
.. warning::
.. warning::
Due to the configuration of the underlying :class:`~dask.array.core.Array`, its :meth:`dask.array.Array.compute` method has to be applied priori to de-dispersing the data.
Therefore, a potential huge data volume may be computed at once.
By default, a security exception is raised to prevent computing a too large data set.
To bypass this limit, set ``'ignore_volume_warning'`` of :attr:`~nenupy.io.tf.TFPipeline.parameters` to `True`.
Due to the configuration of the underlying :class:`~dask.array.core.Array`, its :meth:`dask.array.Array.compute` method has to be applied priori to de-dispersing the data.
Therefore, a potential huge data volume may be computed at once.
By default, a security exception is raised to prevent computing a too large data set.
To bypass this limit, set ``'ignore_volume_warning'`` of :attr:`~nenupy.io.tf.TFPipeline.parameters` to `True`.
"""
def wrapper_task(frequency_hz, data, dt, dispersion_measure, ignore_volume_warning):

def wrapper_task(
frequency_hz, data, dt, dispersion_measure, ignore_volume_warning
):
if dispersion_measure is None:
return frequency_hz, data
return frequency_hz, data
# Make sure the data volume is not too big!
projected_data_volume = data.nbytes * u.byte
if (projected_data_volume >= DATA_VOLUME_SECURITY_THRESHOLD) and (
Expand All @@ -228,7 +231,11 @@ def wrapper_task(frequency_hz, data, dt, dispersion_measure, ignore_volume_warni
)
return frequency_hz, da.from_array(data, chunks=tmp_chuncks)

return cls("De-disperse", wrapper_task, ["dt", "dispersion_measure", "ignore_volume_warning"])
return cls(
"De-disperse",
wrapper_task,
["dt", "dispersion_measure", "ignore_volume_warning"],
)

@classmethod
def time_rebin(cls):
Expand Down Expand Up @@ -338,7 +345,7 @@ def __init__(self, data_obj: Any, *tasks: TFTask):

def __repr__(self) -> str:
return self.info()

@property
def parameters(self) -> utils.TFPipelineParameters:
"""_summary_
Expand All @@ -347,6 +354,7 @@ def parameters(self) -> utils.TFPipelineParameters:
:rtype: :class:`~nenupy.io.tf_utils.TFPipelineParameters`
"""
return self._parameters

@parameters.setter
def parameters(self, params: utils.TFPipelineParameters) -> None:
self._parameters = params
Expand Down Expand Up @@ -545,45 +553,47 @@ def info(self) -> None:
print(message)

def get(self, **pipeline_kwargs) -> SData:
"""_summary_
"""Select time-frequency data, run the user-defined pipeline and return the product.
Data selection, as well as pipeline specific arguments are defined as keyword arguments and passed to :attr:`nenupy.io.tf.TFPipeline.parameters`.
The pipeline can be accessed and modified through the :attr:`nenupy.io.tf.Spectra.pipeline` attribute.
.. rubric:: Available parameters
The following arguments, if set, will update the :class:`~nenupy.io.tf.TFPipeline`'s :attr:`~nenupy.io.tf.TFPipeline.parameters` attribute.
:param tmin: Lower edge of time selection, can either be given as a :class:`~astropy.time.Time` object or an ISOT/ISO string.
:type tmin: `str` or :class:`~astropy.time.Time`
:param tmax: Hello
:param tmax: Upper edge of time selection, can either be given as an :class:`~astropy.time.Time` object or an ISOT/ISO string.
:type tmax: `str` or :class:`~astropy.time.Time`
:param fmin: Hello
:type fmin: `str` or :class:`~astropy.unit.Quantity`
:param fmin: Hello
:type fmax: `str` or :class:`~astropy.unit.Quantity`
:param beam: Hello
:type beam: str or :class:`~astropy.time.Time`
:param dispersion_measure: Hello
:param fmin: Lower frequency boundary selection, can either be given as a :class:`~astropy.unit.Quantity` object or float (assumed to be in MHz in that case).
:type fmin: `float` or :class:`~astropy.unit.Quantity`
:param fmax: Higher frequency boundary selection, can either be given as a :class:`~astropy.unit.Quantity` object or float (assumed to be in MHz in that case).
:type fmax: `float` or :class:`~astropy.unit.Quantity`
:param beam: Beam selection, a single integer corresponding to the index of a recorded numerical beam is expected. Default is the first recorded.
:type beam: `int`
:param dispersion_measure: Enable de-dispersion of the data by this Dispersion Measure. Note that the :meth:`~nenupy.io.tf.TFTask.de_disperse` task should be present in the planned pipeline (:attr:`~nenupy.io.tf.Spectra.pipeline`). It can either be provided as a :class:`~astropy.Quantity` object or a float (assumed to be in pc/cm^3 in that case).
:type dispersion_measure: `float` or :class:`~astropy.unit.Quantity`
:param rotation_measure: Hello
:type rotation_measure: `float` or :class:`~astropy.unit.Quantity`
:param rebin_dt: Hello
:param rebin_dt: Desired rebinning time resolution, can either be given as a :class:`~astropy.unit.Quantity` object or a float (assumed to be in sec in that case). Note that the :meth:`~nenupy.io.tf.TFTask.time_rebin` task should be present in the planned pipeline (:attr:`~nenupy.io.tf.Spectra.pipeline`).
:type rebin_dt: `float` or :class:`~astropy.unit.Quantity`
:param rebin_df: Hello
:param rebin_df: Desired rebinning frequency resolution, can either be given as a :class:`~astropy.unit.Quantity` object or float (assumed to be in kHz in that case). Note that the :meth:`~nenupy.io.tf.TFTask.frequency_rebin` task should be present in the planned pipeline (:attr:`~nenupy.io.tf.Spectra.pipeline`).
:type rebin_df: `float` or :class:`~astropy.unit.Quantity`
:param remove_channels: Hello
:type remove_channels: str or :class:`~astropy.time.Time`
:param dreambeam_skycoord: Hello
:type dreambeam_skycoord: str or :class:`~astropy.time.Time`
:param dreambeam_dt: Hello
:type dreambeam_dt: str or :class:`~astropy.time.Time`
:param dreambeam_parallactic: Hello
:type dreambeam_parallactic: bool
:param stokes: Hello
:type stokes: str or :class:`~astropy.time.Time`
:param ignore_volume_warning: Hello
:type ignore_volume_warning: bool
:return: _description_
:rtype: SData
:param remove_channels: List of subband channels to remove, e.g. `remove_channels=[0,1,-1]` would remove the first, second (low-freq) and last channels from each subband. Note that the :meth:`~nenupy.io.tf.TFTask.remove_channels` task should be present in the planned pipeline (:attr:`~nenupy.io.tf.Spectra.pipeline`).
:type remove_channels: `list` or :class:`~numpy.ndarray`
:param dreambeam_skycoord: Tracked celestial coordinates used during *DreamBeam* correction (along with ``'dreambeam_dt'`` and ``'dreambeam_parallactic'``), a :class:`~astropy.coordinates.SkyCoord` object is expected. Note that the :meth:`~nenupy.io.tf.TFTask.correct_polarization` task should be present in the planned pipeline (:attr:`~nenupy.io.tf.Spectra.pipeline`).
:type dreambeam_skycoord: :class:`~astropy.coordinates.SkyCoord`
:param dreambeam_dt: *DreamBeam* correction time resolution (along with ``'dreambeam_skycoord'`` and ``'dreambeam_parallactic'``), a :class:`~astropy.Quantity` object or a float (assumed in seconds) are expected. Note that the :meth:`~nenupy.io.tf.TFTask.correct_polarization` task should be present in the planned pipeline (:attr:`~nenupy.io.tf.Spectra.pipeline`).
:type dreambeam_dt: `float` or :class:`~astropy.unit.Quantity`
:param dreambeam_parallactic: *DreamBeam* parallactic angle correction (along with ``'dreambeam_skycoord'`` and ``'dreambeam_dt'``), a boolean is expected. Note that the :meth:`~nenupy.io.tf.TFTask.correct_polarization` task should be present in the planned pipeline (:attr:`~nenupy.io.tf.Spectra.pipeline`).
:type dreambeam_parallactic: `bool`
:param stokes: Stokes parameter selection, can either be given as a string or a list of strings, e.g. ['I', 'Q', 'V/I']. Note that the :meth:`~nenupy.io.tf.TFTask.get_stokes` task should be present in the planned pipeline (:attr:`~nenupy.io.tf.Spectra.pipeline`).
:type stokes: `str` or `list[str]`
:param ignore_volume_warning: Ignore or not (default value) the limit regarding output data volume.
:type ignore_volume_warning: `bool`
:return: Processed data selection.
:rtype: :class:`~nenupy.beamlet.sdata.SData`
"""

# Update the pipeline parameters to user's last requests
Expand Down

0 comments on commit 4a459fd

Please sign in to comment.