fix kwave bugs

Author: peterhollender

src/openlifu/plan/solution.py

@@ -256,15 +256,14 @@ def analyze(self,
             apodization = self.apodizations[focus_index]
             origin = self.transducer.get_effective_origin(apodizations=apodization, units=options.distance_units)
 
-            output_signal_Pa = self.transducer.calc_output(
-                input_signal_V,
-                cycles=self.pulse.duration * self.pulse.frequency,
+            p0_Pa = np.max(self.transducer.calc_output(
+                cycles=1,
                 frequency=self.pulse.frequency,
                 dt=dt,
                 delays=self.delays[focus_index, :],
                 apod=self.apodizations[focus_index, :],
-            )
-            p0_Pa = np.max(output_signal_Pa, axis=1)
+                amplitude=self.pulse.amplitude * self.voltage,
+            ), axis=1)
 
             mainlobe_mask = get_mask(
                 pnp_MPa,

src/openlifu/sim/kwave_if.py

@@ -139,7 +139,8 @@ def run_simulation(arr: xdc.Transducer,
     apod = apod if apod is not None else np.ones(arr.numelements())
     kgrid = get_kgrid(params.coords, dt=dt, t_end=t_end, cfl=cfl)
     t = np.arange(0, np.min([cycles / freq, (kgrid.Nt-np.ceil(max(delays)/kgrid.dt))*kgrid.dt]), kgrid.dt)
-    input_signal = amplitude * np.sin(2 * np.pi * freq * t)
+    if cycles/freq > t[-1]:
+        cycles = int(t[-1]*freq)
     units = [params[dim].attrs['units'] for dim in params.dims]
     if not all(unit == units[0] for unit in units):
         raise ValueError("All dimensions must have the same units")
@@ -156,14 +157,15 @@ def run_simulation(arr: xdc.Transducer,
     if _source is not None:
         source = _source
     else:
-        source_mat = arr.calc_output(input_signal, cycles=cycles, frequency=freq, dt=kgrid.dt, delays=delays, apod=apod)
+        source_mat = arr.calc_output(amplitude=amplitude, cycles=cycles, frequency=freq, dt=kgrid.dt, delays=delays, apod=apod)
+        source_mat = source_mat[:, :kgrid.Nt]
     if arr.crosstalk_frac != 0:
         # Simulate crosstalk by adding additional elements to the array for each element that
         #   is within the crosstalk distance, with the signal scaled by the crosstalk fraction.
         #   This is a simple model of crosstalk and may not capture all the complexities of real
         #   crosstalk, but it can be useful for testing and simulation purposes.
-        crosstalk_arr = arr.copy()
-        crosstalk_mat = source_mat
+        #crosstalk_arr = arr.copy()
+        crosstalk_mat = source_mat.copy()
         positions = arr.get_positions(units="m")
         for src_idx in range(arr.numelements()):
             for dst_idx in range(arr.numelements()):
@@ -173,9 +175,10 @@ def run_simulation(arr: xdc.Transducer,
                 dst_pos = np.array(positions[dst_idx])
                 dist = np.linalg.norm(src_pos - dst_pos)
                 if dist <= arr.crosstalk_dist:
-                    crosstalk_arr.elements += [arr.elements[dst_idx].copy()]
-                    crosstalk_mat = np.vstack((crosstalk_mat, arr.crosstalk_frac*source_mat[src_idx,:]))
-        arr = crosstalk_arr
+         #           crosstalk_arr.elements += [arr.elements[dst_idx].copy()]
+                    crosstalk_mat[dst_idx,:] += arr.crosstalk_frac*source_mat[src_idx,:]
+                    #crosstalk_mat = np.vstack((crosstalk_mat, arr.crosstalk_frac*source_mat[src_idx,:]))
+        #arr = crosstalk_arr
         source_mat = crosstalk_mat
     karray = get_karray(arr,
                         translation=array_offset,

src/openlifu/xdc/element.py

@@ -9,16 +9,66 @@
 from openlifu.util.annotations import OpenLIFUFieldData
 from openlifu.util.units import getunitconversion
 
+SENS_FREQ_KEY = "freq_Hz"
+SENS_VALUE_KEY = "values_Pa_per_V"
 
-def sensitivity_at_frequency(sensitivity: float | dict[float, float], frequency: float) -> float:
+
+def normalize_sensitivity(sensitivity: float | dict) -> float | dict[str, list[float]]:
+    """Normalize sensitivity to a canonical representation.
+
+    Canonical frequency-dependent representation is:
+        {"freq_Hz": [...], "values_Pa_per_V": [...]}
+
+    Backward-compatible legacy representation with frequency keys is accepted and
+    converted to the canonical representation.
+    """
+    if isinstance(sensitivity, dict):
+        if SENS_FREQ_KEY in sensitivity or SENS_VALUE_KEY in sensitivity:
+            if SENS_FREQ_KEY not in sensitivity or SENS_VALUE_KEY not in sensitivity:
+                raise ValueError("Sensitivity dictionary must include both 'freq_Hz' and 'values_Pa_per_V'.")
+            freqs = np.asarray(sensitivity[SENS_FREQ_KEY], dtype=np.float64).reshape(-1)
+            values = np.asarray(sensitivity[SENS_VALUE_KEY], dtype=np.float64).reshape(-1)
+        else:
+            # Legacy format: {frequency_hz: sensitivity}
+            if len(sensitivity) == 0:
+                raise ValueError("Sensitivity dictionary must not be empty.")
+            mapping = {float(k): float(v) for k, v in sensitivity.items()}
+            freqs = np.array(list(mapping.keys()), dtype=np.float64)
+            values = np.array(list(mapping.values()), dtype=np.float64)
+
+        if len(freqs) == 0:
+            raise ValueError("Sensitivity frequency list must not be empty.")
+        if len(freqs) != len(values):
+            raise ValueError("Sensitivity frequency and value lists must have the same length.")
+
+        order = np.argsort(freqs)
+        freqs = freqs[order]
+        values = values[order]
+        if np.any(np.diff(freqs) <= 0):
+            raise ValueError("Sensitivity frequencies must be strictly increasing.")
+
+        return {
+            SENS_FREQ_KEY: [float(f) for f in freqs],
+            SENS_VALUE_KEY: [float(v) for v in values],
+        }
+
+    return float(sensitivity)
+
+
+def sensitivity_at_frequency(sensitivity: float | dict, frequency: float) -> float:
+    sensitivity = normalize_sensitivity(sensitivity)
     if isinstance(sensitivity, dict):
-        freqs = np.array(list(sensitivity.keys()), dtype=np.float64)
-        values = np.array(list(sensitivity.values()), dtype=np.float64)
-        return float(np.interp(frequency, freqs, values, left=values[0], right=values[-1]))
+        if frequency in sensitivity[SENS_FREQ_KEY]:
+            idx = sensitivity[SENS_FREQ_KEY].index(frequency)
+            return float(sensitivity[SENS_VALUE_KEY][idx])
+        else:
+            freqs = np.array(sensitivity[SENS_FREQ_KEY], dtype=np.float64)
+            values = np.array(sensitivity[SENS_VALUE_KEY], dtype=np.float64)
+            return float(np.interp(frequency, freqs, values, left=values[0], right=values[-1]))
     return float(sensitivity)
 
 
-def generate_drive_signal(input_signal, cycles: float, frequency: float, dt: float) -> np.ndarray:
+def generate_drive_signal(cycles: float, frequency: float, dt: float, amplitude: float = 1.0) -> np.ndarray:
     """Generate a drive signal with duration constrained by cycles/frequency."""
     if dt <= 0:
         raise ValueError("dt must be positive.")
@@ -27,15 +77,8 @@ def generate_drive_signal(input_signal, cycles: float, frequency: float, dt: flo
     if cycles <= 0:
         raise ValueError("cycles must be positive.")
     n_samples = max(1, int(np.round(cycles / (frequency * dt))))
-    if np.isscalar(input_signal):
-        t = np.arange(n_samples, dtype=np.float64) * dt
-        return float(input_signal) * np.sin(2 * np.pi * frequency * t)
-    base = np.asarray(input_signal, dtype=np.float64).reshape(-1)
-    drive_signal = np.zeros(n_samples, dtype=np.float64)
-    n_copy = min(n_samples, len(base))
-    drive_signal[:n_copy] = base[:n_copy]
-    return drive_signal
-
+    t = np.arange(n_samples, dtype=np.float64) * dt
+    return amplitude * np.sin(2 * np.pi * frequency * t)
 
 def matrix2xyz(matrix):
     x = matrix[0, 3]

src/openlifu/xdc/transducer.py

@@ -112,16 +112,16 @@ def __post_init__(self):
     def get_sensitivity(self, frequency: float) -> float:
         return sensitivity_at_frequency(self.sensitivity, frequency)
 
-    def calc_output(self, input_signal, cycles: float, frequency: float, dt: float, delays: np.ndarray = None, apod: np.ndarray = None):
+    def calc_output(self, cycles: float, frequency: float, dt: float, delays: np.ndarray = None, apod: np.ndarray = None, amplitude: float = 1.0) -> np.ndarray:
         if delays is None:
             delays = np.zeros(self.numelements())
         if apod is None:
             apod = np.ones(self.numelements())
-        drive_signal = generate_drive_signal(input_signal, cycles=cycles, frequency=frequency, dt=dt)
+        drive_signal = generate_drive_signal(cycles=cycles, frequency=frequency, dt=dt, amplitude=amplitude)
         base_output = drive_signal * self.get_sensitivity(frequency)
         outputs = [
             np.concatenate(
-                [np.zeros(int(delay / dt)), a * element.get_sensitivity(frequency) * base_output],
+                [np.zeros(int(delay / dt)),  a * element.get_sensitivity(frequency) * base_output],
                 axis=0,
             )
             for element, delay, a, in zip(self.elements, delays, apod)