fix tests

Author: faridyagubbayli

kwave/data.py

@@ -163,6 +163,35 @@ class SimulationResult:
     I_avg: Optional[np.ndarray] = None
     I: Optional[np.ndarray] = None
 
+    def __getitem__(self, key: str):
+        """
+        Enable dictionary-style access for backward compatibility.
+
+        Args:
+            key: Field name to access
+
+        Returns:
+            Value of the field
+
+        Raises:
+            KeyError: If the field does not exist
+        """
+        if hasattr(self, key):
+            return getattr(self, key)
+        raise KeyError(f"'{key}' field not found in SimulationResult")
+
+    def __contains__(self, key: str) -> bool:
+        """
+        Enable dictionary-style membership testing for backward compatibility.
+
+        Args:
+            key: Field name to check
+
+        Returns:
+            True if the field exists, False otherwise
+        """
+        return hasattr(self, key)
+
     @classmethod
     def from_dotdict(cls, data: dict) -> "SimulationResult":
         """

kwave/kspaceFirstOrder2D.py

@@ -1,4 +1,4 @@
-from typing import Union
+from typing import Optional, Union
 
 import numpy as np
 from beartype import beartype as typechecker
@@ -86,7 +86,7 @@ def kspaceFirstOrder2DC(
     medium: kWaveMedium,
     simulation_options: SimulationOptions,
     execution_options: SimulationExecutionOptions,
-) -> SimulationResult:
+) -> Optional[SimulationResult]:
     """
     2D time-domain simulation of wave propagation using C++ code.
 
@@ -146,7 +146,7 @@ def kspaceFirstOrder2D(
     medium: kWaveMedium,
     simulation_options: SimulationOptions,
     execution_options: SimulationExecutionOptions,
-) -> SimulationResult:
+) -> Optional[SimulationResult]:
     """
     2D time-domain simulation of wave propagation using k-space pseudospectral method.
 

kwave/kspaceFirstOrder3D.py

@@ -1,4 +1,4 @@
-from typing import Union
+from typing import Optional, Union
 
 import numpy as np
 
@@ -26,7 +26,7 @@ def kspaceFirstOrder3DG(
     medium: kWaveMedium,
     simulation_options: SimulationOptions,
     execution_options: SimulationExecutionOptions,
-) -> SimulationResult:
+) -> Optional[SimulationResult]:
     """
     3D time-domain simulation of wave propagation on a GPU using C++ CUDA code.
 
@@ -80,7 +80,7 @@ def kspaceFirstOrder3DC(
     medium: kWaveMedium,
     simulation_options: SimulationOptions,
     execution_options: SimulationExecutionOptions,
-) -> SimulationResult:
+) -> Optional[SimulationResult]:
     """
     3D time-domain simulation of wave propagation using C++ code.
 
@@ -137,7 +137,7 @@ def kspaceFirstOrder3D(
     simulation_options: SimulationOptions,
     execution_options: SimulationExecutionOptions,
     time_rev: bool = False,  # deprecated parameter
-) -> SimulationResult:
+) -> Optional[SimulationResult]:
     """
     3D time-domain simulation of wave propagation using k-space pseudospectral method.
 

kwave/kspaceFirstOrderAS.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Union
+from typing import Optional, Union
 
 import numpy as np
 from numpy.fft import ifftshift
@@ -31,7 +31,7 @@ def kspaceFirstOrderASC(
     medium: kWaveMedium,
     simulation_options: SimulationOptions,
     execution_options: SimulationExecutionOptions,
-) -> SimulationResult:
+) -> Optional[SimulationResult]:
     """
     Axisymmetric time-domain simulation of wave propagation using C++ code.
 
@@ -90,7 +90,7 @@ def kspaceFirstOrderAS(
     medium: kWaveMedium,
     simulation_options: SimulationOptions,
     execution_options: SimulationExecutionOptions,
-) -> SimulationResult:
+) -> Optional[SimulationResult]:
     """
     Axisymmetric time-domain simulation of wave propagation.
 
@@ -159,7 +159,7 @@ def kspaceFirstOrderAS(
     if simulation_options.simulation_type is not SimulationType.AXISYMMETRIC:
         logging.log(
             logging.WARN,
-            "simulation type is not set to axisymmetric while using kSapceFirstOrderAS. " "Setting simulation type to axisymmetric.",
+            "simulation type is not set to axisymmetric while using kSapceFirstOrderAS. Setting simulation type to axisymmetric.",
         )
         simulation_options.simulation_type = SimulationType.AXISYMMETRIC
 
@@ -297,7 +297,7 @@ def kspaceFirstOrderAS(
 
     # option to run simulations without the spatial staggered grid is not
     # supported for the axisymmetric code
-    assert options.use_sg, "Optional input " "UseSG" " is not supported for axisymmetric simulations."
+    assert options.use_sg, "Optional input UseSG is not supported for axisymmetric simulations."
 
     # =========================================================================
     # SAVE DATA TO DISK FOR RUNNING SIMULATION EXTERNAL TO MATLAB