example added

pravirkr · Apr 24, 2024 · 09d4800 · 09d4800
1 parent 114a174
commit 09d4800
Show file tree

Hide file tree

Showing 5 changed files with 236 additions and 1 deletion.
diff --git a/README.md b/README.md
@@ -1,2 +1,30 @@
 # dmt
-Dispersion Measure Transforms
+
+[![GitHub CI](https://github.com/pravirkr/dmt/actions/workflows/ci.yml/badge.svg)](https://github.com/pravirkr/dmt/actions/workflows/ci.yml)
+[![License](https://img.shields.io/github/license/pravirkr/dmt)](https://github.com/pravirkr/dmt/blob/main/LICENSE)
+
+## Dispersion Measure Transforms
+|           |           |
+| --------- | --------- |
+| ![](docs/waterfall.png) | ![](docs/dmt.png) |
+
+
+## Installation
+
+Using [pip](https://pip.pypa.io):
+
+```bash
+pip install -U git+https://github.com/pravirkr/dmt
+```
+
+## Usage
+
+```python
+from dmt.libdmt import FDMT
+
+frb = np.ones((nchans, nsamps), dtype=np.float32)
+thefdmt = FDMT(f_min, f_max, nchans, nsamps, tsamp, dt_max=dt_max, dt_min=0, dt_step=1)
+dmt_transform = thefdmt.execute(frb.astype(np.float32))
+```
+
+
diff --git a/docs/dmt.png b/docs/dmt.png
diff --git a/docs/example.ipynb b/docs/example.ipynb
diff --git a/docs/waterfall.png b/docs/waterfall.png
diff --git a/src/dmt/simulate.py b/src/dmt/simulate.py
@@ -0,0 +1,42 @@
+import numpy as np
+
+
+def get_dmdelays(
+    dm: float,
+    f_min: float,
+    f_max: float,
+    tsamp: float,
+    nchans: int,
+    *,
+    in_samples: bool = True,
+) -> np.ndarray:
+    """Calculate the DM delays for a given frequency range."""
+    foff = (f_max - f_min) / nchans
+    chan_freqs = np.arange(nchans, dtype=np.float64) * foff + f_min
+    delays = dm * 4.148808e3 * (f_min**-2 - chan_freqs**-2)
+    if in_samples:
+        return (delays / tsamp).round().astype(np.int32)
+    return delays
+
+
+def generate_frb(
+    f_min: float,
+    f_max: float,
+    nchans: int,
+    nsamps: int,
+    tsamp: float,
+    dm: float,
+    amp: float = 1,
+    offset: int = 0,
+    width: int = 1,
+    noise_rms: float = 0,
+) -> np.ndarray:
+    """Generate a simple frb signal."""
+    rng = np.random.default_rng()
+    arr = rng.standard_normal((nchans, nsamps)) * noise_rms
+    arr[:, offset : offset + width] += amp
+    delays = get_dmdelays(dm, f_min, f_max, tsamp, nchans, in_samples=True)
+    new_ar = np.zeros_like(arr)
+    for ichan in range(nchans):
+        new_ar[ichan] = np.roll(arr[ichan], -delays[ichan])
+    return new_ar