# Copilot Instructions for WDM Wavelet Package

## Project Overview
High-performance Wilson-Daubechies-Meyer (WDM) wavelet package for gravitational-wave signal processing. Optimized for embedding in JAX/Numba JIT contexts with instance-level backend selection.

---

## Architecture Requirements

### 1. Backend Strategy
- **Primary**: JAX (mandatory, GPU-portable via `jax.device` placement)
- **Secondary**: Numba for CPU initialization (TD filters only, not transforms)
- **Fallback**: Plain Python when optional dependencies absent
- **Pattern**: Pure functions embeddable in `@jax.jit` / `@numba.njit` contexts

### 2. Instance-Level Backend Selection
- Backend choice locked at WDM instantiation: `WDM(..., backend="jax")`
- **NO per-call `backend=` overrides** — eliminate this pattern entirely
- Decided by parent program configuration (set once, immutable)
- Validation occurs in `__post_init__`, not per-method call

### 3. Pure Function & Wrapping Pattern
```python
# Core functions: return raw arrays only (JAX/NumPy)
from wdm_wavelet.core.t2w import t2w_jax
m_L, aligned_len, raw_tf = t2w_jax(M, m_H, signal, filter_taps, MM)

# Class methods: wrap outputs with user-friendly types
wdm = WDM(M=32, K=64, backend="jax")
tf_map = wdm.t2w(strain, sample_rate=4096.0)  # Returns TimeFrequencyMap
```

**Rule**: Core functions (`t2w_jax`, `w2t_jax`, `w2tQ_jax`, etc.) are embeddable; class methods are user-facing.

### 4. Code Quality Standards

#### PEP 8 Compliance
- Parameters: `snake_case` only (no CamelCase)
- Variables: Clear, non-opaque names
- **No builtin shadowing**: `filter` → `filter_taps`, `type` → `dtype_str`, etc.

#### Naming Consistency
- Replace acronyms: `WWS` → `signal`, `m_Layer` → `m_layer`, `BetaOrder` → `beta_order`
- Canonical function names (no aliases): `t2w()` removed; `t2w_jax()` is the API
- Internal functions prefixed with `_`: `_t2w_jax_impl`, `_mirror_extend_jax`

#### Documentation
- **Module docstrings**: Explain public API + paper reference
  ```python
  """
  Forward WDM transform: time series → TF map.
  
  Reference
  ---------
  V. Necula et al., J. Phys.: Conf. Ser. 363, 012032 (2012), Sec. 2
    https://doi.org/10.1088/1742-6596/363/1/012032
  
  Public API
  ----------
  t2w_jax(...) -> raw arrays
  """
  ```
- **NumPy-style docstrings** for all public functions: Parameters, Returns, Notes, Examples
- **Inline comments** explaining non-obvious logic (e.g., why `n_layers = M + 1` is needed)

### 5. Graceful Degradation
- **Guard imports** with try/except + `HAS_*` flags:
  ```python
  try:
      import jax
      HAS_JAX = True
  except ImportError:
      jax = None
      HAS_JAX = False
  ```
- **Fallback decorators** (no-op if dependency absent):
  ```python
  try:
      from numba import njit
  except ImportError:
      def njit(*args, **kwargs):
          return lambda f: f  # No-op decorator
  ```
- Raise informative `ImportError` at use-time if required backend unavailable

### 6. Testing Requirements
- **Backend validation at instantiation** (not per-call):
  ```python
  wdm = WDM(..., backend="invalid")  # Raises at __post_init__
  ```
- **All 20+ unit tests must pass** before merging any changes
- Test suite updated whenever architecture changes
- Run: `pytest test/ -q` or `pytest test/ -v` for details

---

## Development Checklist

Before submitting any PR or feature:

- ✅ Backend locked at instance level (no per-call `backend=` params)
- ✅ Core functions return raw JAX/NumPy arrays only
- ✅ Class methods wrap outputs with TimeFrequencyMap/TimeSeries/etc.
- ✅ PEP 8 compliant (snake_case, no shadowing, clear names)
- ✅ Comprehensive docstrings (NumPy-style + paper refs)
- ✅ Module docstrings with public API overview
- ✅ Optional dependencies guarded (HAS_JAX, HAS_NUMBA checks)
- ✅ All imports inside try/except with fallbacks
- ✅ All 20+ tests passing (`pytest test/ -q`)
- ✅ No redundant aliases (e.g., `t2w()` removed; `t2w_jax()` canonical)

---

## File Structure Reference

### Core Modules
- **`wdm_wavelet/wdm.py`**: Public WDM class + high-level transforms
- **`wdm_wavelet/core/t2w.py`**: Forward transform (time → TF)
- **`wdm_wavelet/core/w2t.py`**: Inverse transforms (TF → time)
- **`wdm_wavelet/core/get_filter.py`**: WDM prototype filter synthesis
- **`wdm_wavelet/core/time_delay.py`**: Delayed WDM amplitudes (TD filters)

### Tests
- **`test/test_wdm.py`**: WDM class tests
- **`test/test_core_wavelet.py`**: Core function tests
- **`test/test_time_delay.py`**: TD filter tests

---

## Common Patterns

### Adding a New Transform Method
```python
def my_transform(self, signal, **kwargs) -> SomeType:
    """
    User-facing method.
    
    Parameters
    ----------
    signal : TimeSeries or array-like
    **kwargs : dict
        Additional options.
    
    Returns
    -------
    SomeType
        Wrapped result.
    """
    # Coerce input
    arr, sr, t0 = self._coerce_timeseries_input(signal, **kwargs)
    
    # Call pure function (embeddable in JIT)
    raw_result = my_transform_jax(self.M, arr, self.filter, ...)
    
    # Wrap result
    return SomeType(raw_result, sr, t0)
```

### Adding a New Core Function
```python
def my_transform_jax(M, signal, filter_taps, other_params):
    """
    Pure function for embedding in @jax.jit contexts.
    
    Parameters
    ----------
    M : int
        Frequency bands.
    signal : array-like
        Input data.
    filter_taps : array-like
        Filter coefficients.
    other_params : type
        Description.
    
    Returns
    -------
    result : jax.Array
        Raw JAX array output.
    
    Notes
    -----
    Safe to embed in @jax.jit / jax.vmap / @numba.njit contexts.
    """
    if not HAS_JAX:
        raise ImportError("JAX required for my_transform_jax")
    
    # Implementation here
    return raw_jax_array
```

### Handling Optional Dependencies
```python
# At module level
try:
    import optional_lib
    HAS_OPTIONAL = True
except ImportError:
    optional_lib = None
    HAS_OPTIONAL = False

# In function
def uses_optional():
    if not HAS_OPTIONAL:
        raise ImportError("optional_lib is required for this feature")
    # Use optional_lib here
```

---

## Performance Notes

- **JAX**: Used for transforms (t2w, w2t). JIT compilation enabled via `@partial(jax.jit, static_argnames=(...))`.
- **Numba**: Used for TD filter inner loops (`_fill_td_filter1_loop`, `_fill_td_filter2_loop`) with `@njit(cache=True)`.
- **Transforms**: NOT embedded in Numba; only JAX embeddable for performance.
- **Filter construction**: Can use Numba or JAX; dispatched in `get_filter()` based on availability.

---

## References

**Paper**: V. Necula et al., J. Phys.: Conf. Ser. 363, 012032 (2012)
https://doi.org/10.1088/1742-6596/363/1/012032

**Key equations**:
- Eq. (19)-(23): Mirror-extended truncated sums (forward transform)
- Time-delay filters: Temporal shifts for coincidence analysis

---

## Questions?

Refer to the conversation summary in the git history or consult the inline docstrings in core modules.
