Metadata-Version: 2.4
Name: pypulse_unal
Version: 0.1.1
Summary: A Python library for pulse analysis.
Home-page: https://gitlab.com/jsierraj/pypulse
Author: Juan Sierra, Juan Ardila
Author-email: juansesierra54@gmail.com
License: : OSI Approved :: MIT License
Requires-Python: >=3.8
Description-Content-Type: text/markdown

# PyPulse

This python library is a JAX-based tool for pulse construction, analysis, and visualization, using JAX for high-performance calculations and GPU acceleration. 
Pypulse provides JAX-compatible functions for constructing electric fields in the frequency domain, applying phase modulation (chirp) and amplitude filtering (notch), performing Fourier transforms to the time domain, and calculating advanced pulse representations like the Wigner function.

## Installation

You can install the library directly from PyPI:

```bash
pip install pypulse_unal
```

-----

## Function Documentation

The functions rely on 2 modules:

  * `pulses`: Provides pulse shape definitions and constants.
  * `aux_functions`: Provides Fourier transform and numerical limit finding utilities.

### `frequency_field`

```python
def frequency_field(
    f: jnp.ndarray,
    f_0: float,
    pulse_fwhm: float,
    pulse_area: float,
    pulse_type: int
) -> jnp.ndarray:
```

Constructs a **pure electric field pulse** in the frequency domain based on a specified shape, bandwidth, and area.

| Type | Name | Description |
| :--- | :--- | :--- |
| **Input** | `f` (`jnp.ndarray`) | Array of **frequencies** where the pulse is evaluated. |
| **Input** | `f_0` (`float`) | **Pulse center frequency** (carrier frequency). |
| **Input** | `pulse_fwhm` (`float`) | **Pulse bandwidth** (Full Width at Half Maximum) in the power profile. |
| **Input** | `pulse_area` (`float`) | Scaling factor for the pulse magnitude, representing the **pulse area**. |
| **Input** | `pulse_type` (`int`) | An integer identifier defining the **shape of the pulse** (e.g., Gaussian, Sech$^2$). |
| **Output** | `jnp.ndarray` | The **evaluated pulse** in the frequency domain, $E(f)$. |

**Operation:**
Calculates $E(f)$ using a shape function $S(f)$ from the external `pulses` module:
$$E(f) = \text{pulse\_area} \cdot S\left(f, f_0, \frac{\text{pulse\_fwhm}}{\text{const}}\right)$$
The shape function's width is normalized using the constant corresponding to `pulse_type` to ensure the bandwidth parameter truly represents the FWHM.

-----

### `set_chirp`

```python
def set_chirp(
    f: jnp.ndarray,
    f_field: jnp.ndarray,
    f_0: float,
    chirp_alpha: float,
    chrip_power: int = 1
) -> jnp.ndarray:
```

Applies a **spectral phase (chirp)** to an existing frequency field. The phase is a polynomial function of the frequency offset $(\omega - \omega_0)$.

| Type | Name | Description |
| :--- | :--- | :--- |
| **Input** | `f` (`jnp.ndarray`) | Array of **frequencies**. |
| **Input** | `f_field` (`jnp.ndarray`) | The un-chirped **frequency field** array, $E(f)$. |
| **Input** | `f_0` (`float`) | **Chirp center frequency** (the frequency around which the phase is zero). |
| **Input** | `chirp_alpha` (`float`) | The **alpha parameter** ($\alpha$) controlling the strength/magnitude of the phase. |
| **Input** | `chirp_power` (`int`) | The **power of the frequency offset** $(\omega - \omega_0)$ in the phase term (e.g., $1$ for linear phase, $2$ for quadratic/GDD, $3$ for cubic/TOD). |
| **Output** | `jnp.ndarray` | The **chirped pulse** in the frequency domain, $E'(f)$. |

**Operation:**
Applies the phase factor $\exp(i \phi(f))$ to the field $E(f)$.
The phase $\phi(f)$ is given by:
$$\phi(f) = \alpha \cdot \frac{\left(2\pi(f-f_0)\right)^{\text{chirp\_power} + 1}}{\text{chirp\_power} + 1}$$
If $\text{chirp\_power} = 0$, a constant phase $\phi(f) = \alpha$ is applied.

-----

### `set_notch`

```python
def set_notch(
    f: jnp.ndarray,
    f_field: jnp.ndarray,
    f_0: float,
    notch_fwhm: float,
    notch_type: int = pulses.gaussian
) -> jnp.ndarray:
```

Applies a **notch filter** (frequency amplitude filter) to the pulse, effectively removing power within a certain frequency band.

| Type | Name | Description |
| :--- | :--- | :--- |
| **Input** | `f` (`jnp.ndarray`) | Array of **frequencies**. |
| **Input** | `f_field` (`jnp.ndarray`) | The **frequency field** array, $E(f)$. |
| **Input** | `f_0` (`float`) | **Notch center frequency**. |
| **Input** | `notch_fwhm` (`float`) | **Notch bandwidth** (FWHM of the removal function). |
| **Input** | `notch_type` (`int`) | An integer identifier defining the **shape of the notch** filter (e.g., Gaussian). |
| **Output** | `jnp.ndarray` | The **filtered pulse** in the frequency domain. |

**Operation:**
The pulse is multiplied by a filter function $F(f)$:
$$E'(f) = E(f) \cdot F(f)$$
where $F(f) = 1 - S(f)$. The function $S(f)$ is the specified shape (`notch_type`) centered at $f_0$ with width $\text{notch\_fwhm}$.

-----

### `time_field`

```python
def time_field(
    f: jnp.ndarray,
    f_field: jnp.ndarray,
    t: jnp.ndarray,
    t_0: float = 0
) -> jnp.ndarray:
```

Calculates the **pulse electric field in the temporal domain**, $E(t)$, by performing the inverse Fourier transform of the frequency field $E(f)$.

| Type | Name | Description |
| :--- | :--- | :--- |
| **Input** | `f` (`jnp.ndarray`) | **Frequencies array** (angular frequencies $\omega = 2\pi f$). |
| **Input** | `f_field` (`jnp.ndarray`) | **Frequency field** array, $E(f)$. |
| **Input** | `t` (`jnp.ndarray`) | **Time array** where the pulse is evaluated. |
| **Input** | `t_0` (`float`) | **Center of the pulse** in time (temporal shift). |
| **Output** | `jnp.ndarray` | The **pulse in the temporal domain**, $E(t)$. |

**Operation:**
Uses the external `aux_functions.vectorized_inverse_fourier` to compute:
$$E(t - t_0) = \int E(f) e^{i 2\pi f (t - t_0)} df$$

-----

### `set_time_limits`

```python
def set_time_limits(
    f: jnp.ndarray,
    f_field: jnp.ndarray,
    step: float,
    t_0: float = 0,
    eps: float = 1e-5,
    max_steps: int = 10_000,
    negative: bool = False
) -> float:
```

Determines the **temporal boundary** (limit) where the pulse's intensity drops below a specified minimum value ($\epsilon$). Useful for setting the boundaries of the time grid.

| Type | Name | Description |
| :--- | :--- | :--- |
| **Input** | `f` (`jnp.ndarray`) | **Frequency domain** array. |
| **Input** | `f_field` (`jnp.ndarray`) | Pulse in **frequency domain**. |
| **Input** | `step` (`float`) | **Size of the steps** used for searching the limit. |
| **Input** | `t_0` (`float`) | **Center of the interval** to search for the limit. |
| **Input** | `eps` (`float`) | **Minimum normalized intensity** value used to set the limit (e.g., $10^{-5}$). |
| **Input** | `max_steps` (`float`) | **Maximum number of iterations** allowed to find the limit. |
| **Input** | `negative` (`bool`) | If `True`, searches for the limit in the negative time direction. |
| **Output** | `float` | The **temporal limit** $t_{lim}$ where $|E(t_{lim})|^2 / \max(|E(t)|^2) < \text{eps}$. |

**Operation:**
This function internally defines the time-domain pulse function, then calls the generic numerical search utility `aux_functions.set_limits` to iteratively find the time coordinate where the pulse intensity falls below `eps`.

-----

### `set_frequency_limits`

```python
def set_frequency_limits(
    f_0: float,
    pulse_fwhm: float,
    pulse_type: int,
    step: float,
    eps: float = 1e-5,
    max_steps: int = 10_000,
    negative: bool = False
) -> float:
```

Determines the **frequency boundary** (limit) where the pulse's intensity drops below a specified minimum value ($\epsilon$). Useful for setting the boundaries of the frequency grid.

| Type | Name | Description |
| :--- | :--- | :--- |
| **Input** | `f_0` (`float`) | **Center of the interval** to search for the limit (pulse center frequency). |
| **Input** | `pulse_fwhm` (`float`) | Pulse bandwidth (used to normalize the pulse width). |
| **Input** | `pulse_type` (`int`) | Shape of the pulse in frequency domain. |
| **Input** | `step` (`float`) | **Size of the steps** used for searching the limit. |
| **Input** | `eps` (`float`) | **Minimum normalized intensity** value used to set the limit. |
| **Input** | `max_steps` (`float`) | **Maximum number of iterations** allowed to find the limit. |
| **Input** | `negative` (`bool`) | If `True`, searches for the limit in the negative frequency direction. |
| **Output** | `float` | The **frequency limit** $f_{lim}$ where $|E(f_{lim})|^2 / \max(|E(f)|^2) < \text{eps}$. |

**Operation:**
This function internally defines the frequency-domain pulse intensity function, then calls the generic numerical search utility `aux_functions.set_limits` to iteratively find the frequency coordinate where the pulse intensity falls below `eps`.

-----

### `get_inst_frequency`

```python
def get_inst_frequency(
        f: jnp.ndarray,
        f_field: jnp.ndarray,
        t: jnp.ndarray,
        t_0: float):
```

Calculates the **instantaneous frequency** $\omega_{inst}(t)$ of the pulse in the temporal domain.

| Type | Name | Description |
| :--- | :--- | :--- |
| **Input** | `f` (`jnp.ndarray`) | **Frequencies array**. |
| **Input** | `f_field` (`jnp.ndarray`) | **Frequency field** array, $E(f)$. |
| **Input** | `t` (`jnp.ndarray`) | **Time array** where the instantaneous frequency is evaluated. |
| **Input** | `t_0` (`float`) | **Temporal shift** of the pulse. |
| **Output** | `jnp.ndarray` | The **instantaneous frequency** $\omega_{inst}(t)$. |

**Operation:**
Calculates the instantaneous frequency using the phase $\phi(t)$ of the temporal electric field $E(t) = |E(t)| e^{i\phi(t)}$:
$$\omega_{inst}(t) = \frac{1}{2\pi} \frac{d\phi(t)}{dt}$$
This calculation is delegated to `aux_functions.aux_instantaneous_frequency`.

-----

### `wigner_function`

```python
def wigner_function(
    t_plot: jnp.ndarray,
    f_plot: jnp.ndarray,
    f_field: jnp.ndarray,
) -> jnp.ndarray:
```

Calculates the **Wigner function** $W(t, f)$, a time-frequency distribution that visualizes the intensity and chirp of the pulse.

| Type | Name | Description |
| :--- | :--- | :--- |
| **Input** | `t_plot` (`jnp.ndarray`) | **Time grid** for the Wigner function evaluation (first output dimension). |
| **Input** | `f_plot` (`jnp.ndarray`) | **Frequency grid** for the Wigner function evaluation (second output dimension). |
| **Input** | `f_field` (`jnp.ndarray`) | **Frequency field** $E(f)$. |
| **Output** | `jnp.ndarray` | The **Wigner function** array $W(t, f)$ with shape $(\text{len}(t), \text{len}(f))$. |

**Operation:**
The Wigner function is calculated using the definition via the temporal field $E(t)$:
$$W(t, f) = \int d\tau' \: E(t + \tau'/2) \cdot E^*(t - \tau'/2) \cdot e^{i 2\pi f \tau'}$$
This calculation is delegated to `aux_functions.compute_wigner`.

-----

You've provided the content for two critical modules, `pulses.py` and `aux_functions.py`, which together handle the fundamental pulse mathematics for your JAX-based library.

Here is the complete documentation for both modules, structured with explanations and detailed input/output tables.


-----

## Module 1: `pulses.py` (Pulse Shape Definitions)

This module defines several common electric field pulse shapes in the frequency domain, normalized to represent their bandwidth based on the **FWHM of the power profile** ($|E(f)|^2$).

### 1\. Curve Shape Definitions

#### `gaussian_curve`

```python
@jit
def gaussian_curve(
    f: jnp.ndarray,
    f_0: float,
    gamma_0: float
) -> jnp.ndarray:
```

Defines a **Gaussian** pulse shape $E(f)$.

| Type | Name | Description |
| :--- | :--- | :--- |
| **Input** | `f` (`jnp.ndarray`) | Frequencies array [PHz]. |
| **Input** | `f_0` (`float`) | Pulse center frequency [PHz]. |
| **Input** | `gamma_0` (`float`) | Pulse bandwidth parameter related to the FWHM of $|E(f)|^2$. |
| **Output** | `jnp.ndarray` | Evaluated electric field in the frequency domain. |

**Operation:**
$$E(f) = \exp\left(-\frac{1}{2}\left(\frac{f-f_0}{\gamma_0}\right)^2\right)$$

#### `lorentzian_curve`

```python
@jit
def lorentzian_curve(
    f: jnp.ndarray,
    f_0: float,
    gamma_0: float
) -> jnp.ndarray:
```

Defines a **Lorentzian** pulse shape $E(f)$.

| Type | Name | Description |
| :--- | :--- | :--- |
| **Input** | `f` (`jnp.ndarray`) | Frequencies array [PHz]. |
| **Input** | `f_0` (`float`) | Pulse center frequency [PHz]. |
| **Input** | `gamma_0` (`float`) | Pulse bandwidth parameter related to the FWHM of $|E(f)|^2$. |
| **Output** | `jnp.ndarray` | Evaluated electric field in the frequency domain. |

**Operation:**
$$E(f) = \frac{1}{1 + \left(\frac{f-f_0}{\gamma_0}\right)^2}$$

#### `psquare_curve`

```python
@jit
def psquare_curve(
    f: jnp.ndarray,
    f_0: float,
    gamma_0: float,
    smoothness: float = 0.005
) -> jnp.ndarray:
```

Defines a **Smoothed Square (or Super-Gaussian)** pulse shape $E(f)$ using a pair of sigmoid functions.

| Type | Name | Description |
| :--- | :--- | :--- |
| **Input** | `f` (`jnp.ndarray`) | Frequencies array [PHz]. |
| **Input** | `f_0` (`float`) | Pulse center frequency [PHz]. |
| **Input** | `gamma_0` (`float`) | Half-width parameter related to the FWHM of $|E(f)|^2$. |
| **Input** | `smoothness` (`float`) | Parameter controlling the sharpness of the edges (lower is sharper). |
| **Output** | `jnp.ndarray` | Evaluated electric field in the frequency domain. |

**Operation:**
The shape is defined by the product of a rising sigmoid (left edge) and a falling sigmoid (right edge) centered at $f_0 \pm \gamma_0$.

#### `sech_curve`

```python
@jit
def sech_curve(
    f: jnp.ndarray,
    f_0: float,
    gamma_0: float
) -> jnp.ndarray:
```

Defines a **Sech** pulse shape $E(f)$.

| Type | Name | Description |
| :--- | :--- | :--- |
| **Input** | `f` (`jnp.ndarray`) | Frequencies array [PHz]. |
| **Input** | `f_0` (`float`) | Pulse center frequency [PHz]. |
| **Input** | `gamma_0` (`float`) | Pulse bandwidth parameter related to the FWHM of $|E(f)|^2$. |
| **Output** | `jnp.ndarray` | Evaluated electric field in the frequency domain. |

**Operation:**
$$E(f) = \text{sech}\left(\frac{f-f_0}{\gamma_0}\right) = \frac{1}{\cosh\left(\frac{f-f_0}{\gamma_0}\right)}$$

#### `sech2_curve`

```python
@jit
def sech2_curve(
    f: jnp.ndarray,
    f_0: float,
    gamma_0: float
) -> jnp.ndarray:
```

Defines a **Sech squared ($\text{Sech}^2$)** pulse shape $E(f)$.

| Type | Name | Description |
| :--- | :--- | :--- |
| **Input** | `f` (`jnp.ndarray`) | Frequencies array [PHz]. |
| **Input** | `f_0` (`float`) | Pulse center frequency [PHz]. |
| **Input** | `gamma_0` (`float`) | Pulse bandwidth parameter related to the FWHM of $|E(f)|^2$. |
| **Output** | `jnp.ndarray` | Evaluated electric field in the frequency domain. |

**Operation:**
$$E(f) = \text{sech}^2\left(\frac{f-f_0}{\gamma_0}\right) = \frac{1}{\cosh^2\left(\frac{f-f_0}{\gamma_0}\right)}$$

### 2\. Constants and Lookups

| Name | Type | Description |
| :--- | :--- | :--- |
| `gaussian` | `int` (0) | Index for `gaussian_curve`. |
| `lorentzian` | `int` (1) | Index for `lorentzian_curve`. |
| `square` | `int` (2) | Index for `psquare_curve`. |
| `sech` | `int` (3) | Index for `sech_curve`. |
| `sech2` | `int` (4) | Index for `sech2_curve`. |
| `curve_shapes` | `list` | List mapping curve index to the JAX function. |
| `fmwh_const` | `jnp.ndarray` | Constants used to convert the input `gamma_0` parameter into the exact FWHM of the power profile ($|E(f)|^2$). |

-----

-----

## Module 2: `aux_functions.py` (Mathematical Utilities)

This module provides utilities for numerical integration, vectorized operations, limit finding, and time-frequency analysis, essential for transforming and analyzing the electric field.

### 1\. Fourier Transforms

#### `scalar_inverse_fourier`

```python
@partial(jit) 
def scalar_inverse_fourier(
    f: jnp.ndarray,
    f_field: jnp.ndarray,
    t: float,
) -> float:
```

Calculates the **inverse Fourier transform** for a single point in time, $E(t)$. Used as the base function for the vectorized version.

| Type | Name | Description |
| :--- | :--- | :--- |
| **Input** | `f` (`jnp.ndarray`) | **Frequencies array**. |
| **Input** | `f_field` (`jnp.ndarray`) | Pulse in the **frequency domain**, $E(f)$. |
| **Input** | `t` (`float`) | **Scalar point in time** to evaluate the field. |
| **Output** | `float` | Evaluated **pulse in the time domain**, $E(t)$. |

**Operation:**
Performs a numerical Riemann sum approximation of the inverse Fourier integral:
$$E(t) \approx \sum_f E(f) e^{-i 2\pi f t} \Delta f$$

#### `vectorized_inverse_fourier`

```python
vectorized_inverse_fourier = vmap(
    scalar_inverse_fourier, 
    in_axes=(None, None, 0)
)
```

A **vectorized version** of `scalar_inverse_fourier` using JAX's `vmap` to efficiently compute $E(t)$ over an array of time points.

### 2\. Instantaneous Frequency

#### `scalar_instantaneous_frequency`

```python
@partial(jit)
def scalar_instantaneous_frequency(
    f: jnp.ndarray,
    f_field: jnp.ndarray,
    t: float,
) -> float:
```

Calculates the **instantaneous frequency** $\omega_{inst}(t)$ for a single time point.

| Type | Name | Description |
| :--- | :--- | :--- |
| **Input** | `f` (`jnp.ndarray`) | **Frequencies array**. |
| **Input** | `f_field` (`jnp.ndarray`) | Pulse in the **frequency domain**, $E(f)$. |
| **Input** | `t` (`float`) | **Scalar point in time** to evaluate the field. |
| **Output** | `float` | **Instantaneous frequency** $\omega_{inst}(t)$ (in Hz). |

**Operation:**
Uses the definition based on the temporal field $E(t)$ and its time derivative $\partial_t E(t)$:
$$\omega_{inst}(t) = -\frac{1}{2\pi} \text{Im}\left\{\frac{\partial_t E(t)}{E(t)}\right\}$$
The derivative $\partial_t E(t)$ is calculated in the frequency domain as the inverse Fourier transform of $-i 2\pi f E(f)$.

#### `aux_instantaneous_frequency`

```python
aux_instantaneous_frequency = vmap(
    scalar_instantaneous_frequency,
    in_axes=(None, None, 0)
)
```

A **vectorized version** of `scalar_instantaneous_frequency` using JAX's `vmap` to compute the instantaneous frequency over an array of time points.

### 3\. Numerical Limit Finding

#### `set_limits`

```python
def set_limits(
    func: Callable,
    step: float,
    f_0: float,
    eps: float = 1e-5,
    max_steps = 10_000,
    negative: bool = False,
) -> float:
```

A generic numerical function to find the coordinate limit (time or frequency) where the **normalized absolute value of a function** (e.g., pulse field) drops below a specified tolerance ($\epsilon$).

| Type | Name | Description |
| :--- | :--- | :--- |
| **Input** | `func` (`Callable`) | The **function** $F(x)$ to be evaluated (e.g., $|E(t)|$). |
| **Input** | `step` (`float`) | **Size of the steps** taken during the search. |
| **Input** | `f_0` (`float`) | **Center of the interval** to start the search. |
| **Input** | `eps` (`float`) | **Minimum normalized value** to set the limit (e.g., $10^{-5}$ of the peak). |
| **Input** | `max_steps` (`float`) | **Maximum number of iterations** allowed. |
| **Input** | `negative` (`bool`) | If `True`, searches in the negative direction from $f_0$. |
| **Output** | `float` | The **coordinate limit** (distance from $f_0$) found, or the maximum distance if limit is not reached. |

**Operation:**
The function iteratively evaluates the ratio $|F(x)|/|F(f_0)|$ at increasing steps away from $f_0$. The process stops when the ratio is less than $\epsilon$ or $\text{max\_steps}$ is reached.

### 4\. Time-Frequency Analysis

#### `compute_wigner`

```python
def compute_wigner(
    t_plot: jnp.ndarray,
    f_plot: jnp.ndarray,
    f_arr: jnp.ndarray,
    f_field: jnp.ndarray,
) -> jnp.ndarray:
```

Computes the **Wigner function** $W(t, f)$, a time-frequency phase space representation, using a memory-efficient `lax.fori_loop` structure vectorized over the 2D grid.

| Type | Name | Description |
| :--- | :--- | :--- |
| **Input** | `t_plot` (`jnp.ndarray`) | **Time grid** for the Wigner function (first output dimension). |
| **Input** | `f_plot` (`jnp.ndarray`) | **Frequency grid** for the Wigner function (second output dimension). |
| **Input** | `f_arr` (`jnp.ndarray`) | **Frequency array** used for the inverse Fourier transforms (the grid of $E(f)$). |
| **Input** | `f_field` (`jnp.ndarray`) | The **field in the frequency domain**, $E(f)$. |
| **Output** | `jnp.ndarray` | **Wigner function** array $W(t, f)$ with shape $(\text{len}(t), \text{len}(f))$. |

**Operation:**
Calculates the integral definition of the Wigner function $W(t, f)$:
$$W(t, f) = \int d\tau \: E(t + \tau/2) \cdot E^*(t - \tau/2) \cdot e^{i 2\pi f \tau}$$
The computation is implemented using nested `vmap` and `lax.fori_loop` to perform the $\tau$ integration at each $(t, f)$ point efficiently.
