eztaox.ts_utils

eztaox.ts_utils#

Utility functions for time series processing.

Functions#

down_sample_by_time(→ tinygp.helpers.JAXArray)

Downsample t_in to match the time points in t_out.

formatlc(→ tuple[tuple[tinygp.helpers.JAXArray, ...)

Transform light curves in dictionary to EzTaoX friendly format.

add_noise(→ tinygp.helpers.JAXArray)

Add Gaussian noise to a time series given measurement uncertainties.

merge_two_sorted_argsort(a, b)

Merge two sorted arrays and return the argsort permutation.

merge_many_sorted_argsort(arrays)

Merge multiple sorted arrays and return the argsort permutation.

merge_sort(→ tinygp.helpers.JAXArray)

Merge multiple sorted arrays and return the argsort permutation.

Module Contents#

down_sample_by_time(t_in, t_out) → tinygp.helpers.JAXArray[source]#

Downsample t_in to match the time points in t_out.

Parameters:

  • t_in (JAXArray) – Array of time values to be downsampled.

  • t_out (JAXArray) – Array of target time values.

Returns:

Downsampled array of time values.

Return type:

JAXArray

formatlc(ts: dict[str, numpy.typing.NDArray | tinygp.helpers.JAXArray], ys: dict[str, numpy.typing.NDArray | tinygp.helpers.JAXArray], yerrs: dict[str, numpy.typing.NDArray | tinygp.helpers.JAXArray], band_order: dict[str, int]) → tuple[tuple[tinygp.helpers.JAXArray, tinygp.helpers.JAXArray], tinygp.helpers.JAXArray, tinygp.helpers.JAXArray][source]#

Transform light curves in dictionary to EzTaoX friendly format.

Parameters:

  • ts (dict[str, NDArray | JAXArray]) – Time stamps for observation in each band.

  • ys (dict[str, NDArray | JAXArray]) – Observed values in each band.

  • yerrs (dict[str, NDArray | JAXArray]) – Uncertainties in observed values for each band.

  • band_order (dict[str, int]) – Mapping of band names to band indices.

Returns:

Light curves formatted as

((time stamps, band indices), observed values, uncertainties).

Return type:

tuple[tuple[JAXArray, JAXArray], JAXArray, JAXArray]

add_noise(y: tinygp.helpers.JAXArray, yerr: tinygp.helpers.JAXArray, key: jax.random.PRNGKey) → tinygp.helpers.JAXArray[source]#

Add Gaussian noise to a time series given measurement uncertainties.

JAX-compatible (works with jit/vmap).

Parameters:

  • y (JAXArray) – Input values to which noise will be added.

  • yerr (JAXArray) – Associated errors for the input values.

  • key (jax.random.PRNGKey) – Pseudorandom number generator key used to draw the noise samples.

Returns:

Array of the same shape as y with additive Gaussian noise applied.

Return type:

JAXArray

merge_two_sorted_argsort(a, b)[source]#

Merge two sorted arrays and return the argsort permutation.

Parameters:

  • a (JAXArray) – First sorted array.

  • b (JAXArray) – Second sorted array.

Returns:

Indices that would sort the concatenation of a and b.

Return type:

JAXArray

merge_many_sorted_argsort(arrays)[source]#

Merge multiple sorted arrays and return the argsort permutation.

Parameters:

arrays (list[JAXArray]) – List of sorted arrays to merge.

Returns:

Indices that would sort the concatenation of all input arrays.

Return type:

JAXArray

merge_sort(*arrays) → tinygp.helpers.JAXArray[source]#

Merge multiple sorted arrays and return the argsort permutation.

Parameters:

*arrays – Variable number of sorted arrays to merge.

Returns:

Indices that would sort the concatenation of all input arrays.

Return type:

JAXArray