project/src/features/engineering.py

"""
Physics-informed feature engineering for femtosecond laser processing.

Derives key physical quantities from raw laser parameters following
established laser-material interaction physics:
- Pulse energy from average power and repetition rate
- Peak fluence from pulse energy and beam geometry (Gaussian beam)
- Effective pulse number from scanning parameters
- Spatial overlap from scan speed and spot size

References
----------
- Perry et al., J. Biomed. Opt. (2009) - fs-laser ablation of hydrogels
- Gamaly et al., Physics Reports (2002) - fs-laser ablation theory
- Petit et al., JLMN (2025) - physics-informed ML features for laser machining
"""

from __future__ import annotations

import logging
from typing import Optional

import numpy as np
import pandas as pd

logger = logging.getLogger(__name__)


def compute_pulse_energy(power_mW: np.ndarray, rep_rate_kHz: np.ndarray) -> np.ndarray:
    """
    Compute pulse energy from average power and repetition rate.

    E_pulse = P_avg / f_rep

    Parameters
    ----------
    power_mW : array-like
        Average laser power in milliwatts
    rep_rate_kHz : array-like
        Pulse repetition rate in kilohertz

    Returns
    -------
    np.ndarray
        Pulse energy in microjoules (µJ)
    """
    return np.asarray(power_mW) / np.maximum(np.asarray(rep_rate_kHz), 1e-6)


def compute_peak_fluence(
    pulse_energy_uJ: np.ndarray,
    spot_diameter_um: np.ndarray,
) -> np.ndarray:
    """
    Compute peak fluence for a Gaussian beam profile.

    F_peak = 2 * E_pulse / (π * w₀²)

    where w₀ = spot_diameter / 2 is the beam waist radius.

    Parameters
    ----------
    pulse_energy_uJ : array-like
        Pulse energy in microjoules
    spot_diameter_um : array-like
        Beam spot diameter (1/e²) in micrometers

    Returns
    -------
    np.ndarray
        Peak fluence in J/cm²
    """
    E_J = np.asarray(pulse_energy_uJ) * 1e-6  # µJ → J
    r_cm = np.asarray(spot_diameter_um) / 2 * 1e-4  # µm → cm (radius)
    return 2 * E_J / (np.pi * np.maximum(r_cm, 1e-8) ** 2)


def compute_effective_pulses(
    rep_rate_kHz: np.ndarray,
    spot_diameter_um: np.ndarray,
    scan_speed_mm_s: np.ndarray,
) -> np.ndarray:
    """
    Compute effective number of pulses per spot.

    N_eff = f_rep × d_spot / v_scan

    This represents the average number of laser pulses that irradiate
    each point along the scan path.

    Parameters
    ----------
    rep_rate_kHz : array-like
        Repetition rate in kHz
    spot_diameter_um : array-like
        Spot diameter in µm
    scan_speed_mm_s : array-like
        Scanning speed in mm/s

    Returns
    -------
    np.ndarray
        Effective pulse count (dimensionless)
    """
    rep_Hz = np.asarray(rep_rate_kHz) * 1000
    spot_mm = np.asarray(spot_diameter_um) * 1e-3
    speed = np.maximum(np.asarray(scan_speed_mm_s), 1e-6)
    return rep_Hz * spot_mm / speed


def compute_spatial_overlap(
    scan_speed_mm_s: np.ndarray,
    rep_rate_kHz: np.ndarray,
    spot_diameter_um: np.ndarray,
) -> np.ndarray:
    """
    Compute spatial pulse overlap percentage.

    Overlap = (1 - pitch / d_spot) × 100%

    where pitch = v_scan / f_rep is the distance between consecutive pulses.

    Parameters
    ----------
    scan_speed_mm_s : array-like
    rep_rate_kHz : array-like
    spot_diameter_um : array-like

    Returns
    -------
    np.ndarray
        Overlap percentage [0, 100)
    """
    pitch_um = np.asarray(scan_speed_mm_s) * 1000 / (np.asarray(rep_rate_kHz) * 1000)
    # pitch in µm: (mm/s * 1000 µm/mm) / (kHz * 1000 Hz/kHz) = µm
    # Actually: pitch = v/f = (mm/s) / (kHz*1000) = mm, convert to µm
    pitch_um = np.asarray(scan_speed_mm_s) / (np.asarray(rep_rate_kHz) * 1000) * 1000  # mm→µm
    spot = np.maximum(np.asarray(spot_diameter_um), 1e-6)
    overlap = np.clip(1 - pitch_um / spot, 0, 0.999)
    return overlap * 100


def compute_energy_dose(
    peak_fluence: np.ndarray,
    effective_pulses: np.ndarray,
) -> np.ndarray:
    """
    Compute total energy dose (accumulated fluence).

    Dose = F_peak × N_eff

    Parameters
    ----------
    peak_fluence : array-like
        Peak fluence in J/cm²
    effective_pulses : array-like
        Effective pulse number

    Returns
    -------
    np.ndarray
        Energy dose in J/cm²
    """
    return np.asarray(peak_fluence) * np.asarray(effective_pulses)


def compute_all_derived_features(df: pd.DataFrame) -> pd.DataFrame:
    """
    Compute all physics-informed derived features and add to DataFrame.

    Parameters
    ----------
    df : pd.DataFrame
        DataFrame with raw laser parameters

    Returns
    -------
    pd.DataFrame
        DataFrame with added derived feature columns
    """
    df = df.copy()

    # Pulse energy
    if "pulse_energy_uJ" not in df.columns:
        df["pulse_energy_uJ"] = compute_pulse_energy(
            df["power_mW"].values, df["repetition_rate_kHz"].values
        )
        logger.info("Computed: pulse_energy_uJ")

    # Peak fluence
    if "peak_fluence_J_cm2" not in df.columns:
        df["peak_fluence_J_cm2"] = compute_peak_fluence(
            df["pulse_energy_uJ"].values, df["spot_diameter_um"].values
        )
        logger.info("Computed: peak_fluence_J_cm2")

    # Effective pulses
    if "effective_pulses" not in df.columns:
        df["effective_pulses"] = compute_effective_pulses(
            df["repetition_rate_kHz"].values,
            df["spot_diameter_um"].values,
            df["scan_speed_mm_s"].values,
        )
        logger.info("Computed: effective_pulses")

    # Overlap
    if "overlap_percent" not in df.columns:
        df["overlap_percent"] = compute_spatial_overlap(
            df["scan_speed_mm_s"].values,
            df["repetition_rate_kHz"].values,
            df["spot_diameter_um"].values,
        )
        logger.info("Computed: overlap_percent")

    # Energy dose
    if "energy_dose_J_cm2" not in df.columns:
        df["energy_dose_J_cm2"] = compute_energy_dose(
            df["peak_fluence_J_cm2"].values, df["effective_pulses"].values
        )
        logger.info("Computed: energy_dose_J_cm2")

    return df


def get_feature_groups(config: dict) -> dict:
    """
    Get feature groups from configuration for ablation studies.

    Parameters
    ----------
    config : dict
        Experiment configuration dictionary

    Returns
    -------
    dict
        Mapping from group name to list of feature column names
    """
    data_cfg = config["data"]
    return {
        "laser": data_cfg["laser_features"],
        "material": data_cfg["material_features"],
        "derived": data_cfg["derived_features"],
        "all": (
            data_cfg["laser_features"]
            + data_cfg["material_features"]
            + data_cfg["derived_features"]
        ),
    }