Merge pull request #67 from symbiotic-engineering/update-notebook

capture excitation force/phase and fix type hints

Author: HopeBestWorld

package/src/openflash/meem_engine.py

@@ -587,6 +587,11 @@ def run_and_store_results(self, problem_index: int) -> Results:
 
         full_added_mass_matrix = np.full((num_freqs, num_modes, num_modes), np.nan)
         full_damping_matrix = np.full((num_freqs, num_modes, num_modes), np.nan)
+        # --- NEW: Initialize matrices for Excitation Force and Phase ---
+        # Note: These are vectors (one value per mode per frequency), not NxN matrices like added mass
+        full_excitation_force = np.full((num_freqs, num_modes), np.nan)
+        full_excitation_phase = np.full((num_freqs, num_modes), np.nan)
+        # ---------------------------------------------------------------
         all_potentials_batch_data = []
 
         # 1. Create a SINGLE reusable Geometry/Problem Setup
@@ -684,6 +689,13 @@ def run_and_store_results(self, problem_index: int) -> Results:
                             j_idx = j_idx_result[0]
                             full_added_mass_matrix[freq_idx, j_idx, i_idx] = coeff_dict['real']
                             full_damping_matrix[freq_idx, j_idx, i_idx] = coeff_dict['imag']
+                            # --- NEW: Capture Excitation Force & Phase ---
+                            # We only need to store this when i_idx (radiating mode) == j_idx (force mode)
+                            # or just capture it for the active mode.
+                            if i_idx == j_idx:
+                                full_excitation_force[freq_idx, j_idx] = coeff_dict.get('excitation_force', np.nan)
+                                full_excitation_phase[freq_idx, j_idx] = coeff_dict.get('excitation_phase', np.nan)
+                            # ---------------------------------------------
 
                     Cs = temp_engine.reformat_coeffs(X_i, NMK_list, len(NMK_list) - 1)
 
@@ -710,10 +722,13 @@ def run_and_store_results(self, problem_index: int) -> Results:
                     full_damping_matrix[freq_idx, :, i_idx] = np.nan
                     continue 
 
+        # --- UPDATED: Pass the new matrices to store_hydrodynamic_coefficients ---
         results.store_hydrodynamic_coefficients(
             frequencies=omegas_to_run,
             added_mass_matrix=full_added_mass_matrix,
             damping_matrix=full_damping_matrix,
+            excitation_force=full_excitation_force, # <--- Add this
+            excitation_phase=full_excitation_phase  # <--- Add this
         )
 
         if all_potentials_batch_data:

package/src/openflash/problem_cache.py

@@ -1,6 +1,6 @@
 # package/src/openflash/problem_cache.py
 import numpy as np
-from typing import Callable, Dict, Any, Optional
+from typing import Callable, Dict, Any, Optional, List
 
 from openflash.multi_equations import *
 
@@ -21,7 +21,7 @@ def __init__(self, problem):
         self.cached_m0: Optional[float] = None 
         # -----------------------------------------------------------------
 
-        self.I_nm_vals: Optional[np.ndarray] = None
+        self.I_nm_vals: Optional[List[np.ndarray]] = None
         self.named_closures: Dict[str, Any] = {}
 
     def _set_A_template(self, A_template: np.ndarray):
@@ -50,7 +50,7 @@ def _set_precomputed_m_k_N_k(self, m_k_arr: np.ndarray, N_k_arr: np.ndarray, m0:
         self.cached_m0 = m0
     # ------------------------------------------------
 
-    def _set_I_nm_vals(self, I_nm_vals: np.ndarray):
+    def _set_I_nm_vals(self, I_nm_vals: List[np.ndarray]):
         self.I_nm_vals = I_nm_vals
 
     def _get_A_template(self) -> np.ndarray:

package/src/openflash/results.py

@@ -1,3 +1,4 @@
+from typing import Optional
 import xarray as xr
 import numpy as np
 from openflash.geometry import Geometry
@@ -203,13 +204,9 @@ def store_all_potentials(self, all_potentials_batch: list[dict]):
 
 
     def store_hydrodynamic_coefficients(self, frequencies: np.ndarray,
-                                        added_mass_matrix: np.ndarray, damping_matrix: np.ndarray):
+                                        added_mass_matrix: np.ndarray, damping_matrix: np.ndarray, excitation_force: Optional[np.ndarray] = None, excitation_phase: Optional[np.ndarray] = None):
         """
-        Store hydrodynamic coefficients (added mass and damping).
-
-        :param frequencies: Array of frequency values.
-        :param added_mass_matrix: 3D array (frequencies x modes x modes) of added mass coefficients.
-        :param damping_matrix: 3D array (frequencies x modes x modes) of damping coefficients.
+        Store hydrodynamic coefficients (added mass, damping, excitation).
         """
         expected_shape = (len(frequencies), len(self.modes), len(self.modes))
 
@@ -222,6 +219,18 @@ def store_hydrodynamic_coefficients(self, frequencies: np.ndarray,
 
         self.dataset['added_mass'] = (('frequency', 'mode_i', 'mode_j'), added_mass_matrix)
         self.dataset['damping'] = (('frequency', 'mode_i', 'mode_j'), damping_matrix)
+        
+        # --- NEW: Store Excitation Data ---
+        if excitation_force is not None:
+            # Shape is (frequency, mode_i) - Force on body i due to its own motion/wave
+            # Note: For diffraction problem strictly, this might differ, but for radiation 
+            # we align it with mode_i.
+            self.dataset['excitation_force'] = (('frequency', 'mode_i'), excitation_force)
+            
+        if excitation_phase is not None:
+            self.dataset['excitation_phase'] = (('frequency', 'mode_i'), excitation_phase)
+        # ----------------------------------
+        
         print("Hydrodynamic coefficients stored in xarray dataset.")
 
     def export_to_netcdf(self, file_path: str):