Coverage for kwave/executor.py: 16%

1import os 

2import stat 

3import subprocess 

4import sys 

5 

6import h5py 

7 

8import kwave 

9from kwave.options.simulation_execution_options import SimulationExecutionOptions 

10from kwave.utils.dotdictionary import dotdict 

11 

12 

13class Executor: 

14 def __init__(self, execution_options: SimulationExecutionOptions, simulation_options): 

15 self.execution_options = execution_options 

16 self.simulation_options = simulation_options 

17 

18 if os.environ.get("KWAVE_FORCE_CPU") == "1": 

19 self.execution_options.is_gpu_simulation = False 

20 self.execution_options.binary_name = "kspaceFirstOrder-OMP" 

21 self.execution_options.binary_path = kwave.BINARY_PATH / self.execution_options.binary_name 

22 self._make_binary_executable() 

23 

24 def _make_binary_executable(self): 

25 binary_path = self.execution_options.binary_path 

26 if not binary_path.exists(): 

27 if kwave.PLATFORM == "darwin" and self.execution_options.is_gpu_simulation: 

28 raise ValueError( 

29 "GPU simulations are currently not supported on MacOS. " 

30 "Try running the simulation on CPU by setting is_gpu_simulation=False." 

31 ) 

32 raise FileNotFoundError(f"Binary not found at {binary_path}") 

33 binary_path.chmod(binary_path.stat().st_mode | stat.S_IEXEC) 

34 

35 def run_simulation(self, input_filename: str, output_filename: str, options: list[str]) -> dotdict: 

36 command = [str(self.execution_options.binary_path), "-i", input_filename, "-o", output_filename] + options 

37 

38 try: 

39 with subprocess.Popen( 

40 command, env=self.execution_options.env_vars, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True 

41 ) as proc: 

42 stdout, stderr = "", "" 

43 if self.execution_options.show_sim_log: 

44 # Stream stdout in real-time 

45 for line in proc.stdout: 

46 print(line, end="") 

47 

48 stdout, stderr = proc.communicate() 

49 

50 proc.wait() # wait for process to finish before checking return code 

51 if proc.returncode != 0: 

52 raise subprocess.CalledProcessError(proc.returncode, command, stdout, stderr) 

53 

54 except subprocess.CalledProcessError as e: 

55 # This ensures stdout is printed regardless of show_sim_logs value if an error occurs 

56 print(e.stdout) 

57 print(e.stderr, file=sys.stderr) 

58 raise 

59 

60 sensor_data = self.parse_executable_output(output_filename) 

61 if not self.simulation_options.pml_inside: 

62 self._crop_pml(sensor_data) 

63 

64 return sensor_data 

65 

66 def _crop_pml(self, sensor_data: dotdict): 

67 Nx = sensor_data["Nx"].item() 

68 Ny = sensor_data["Ny"].item() 

69 Nz = sensor_data["Nz"].item() 

70 pml_x_size = sensor_data["pml_x_size"].item() 

71 pml_y_size = sensor_data["pml_y_size"].item() 

72 pml_z_size = 0 if Nz <= 1 else sensor_data["pml_z_size"].item() 

73 axisymmetric = sensor_data["axisymmetric_flag"].item() 

74 

75 # if the PML is outside, set the index variables to remove the pml 

76 # from the _all and _final variables 

77 x1 = pml_x_size 

78 x2 = Nx - pml_x_size 

79 y1 = 0 if axisymmetric else pml_y_size 

80 y2 = Ny - pml_y_size 

81 z1 = pml_z_size 

82 z2 = Nz - pml_z_size 

83 

84 possible_fields = [ 

85 "p_final", 

86 "p_max_all", 

87 "p_min_all", 

88 "ux_max_all", 

89 "uy_max_all", 

90 "uz_max_all", 

91 "ux_min_all", 

92 "uy_min_all", 

93 "uz_min_all", 

94 "ux_final", 

95 "uy_final", 

96 "uz_final", 

97 ] 

98 for field in possible_fields: 

99 if field in sensor_data: 

100 if sensor_data[field].ndim == 2: 

101 sensor_data[field] = sensor_data[field][y1:y2, x1:x2] 

102 else: 

103 sensor_data[field] = sensor_data[field][z1:z2, y1:y2, x1:x2] 

104 

105 @staticmethod 

106 def parse_executable_output(output_filename: str) -> dotdict: 

107 # Load the C++ data back from disk using h5py 

108 with h5py.File(output_filename, "r") as output_file: 

109 sensor_data = {} 

110 for key in output_file.keys(): 

111 sensor_data[key] = output_file[f"/{key}"][:].squeeze() 

112 

113 # if self.simulation_options.cuboid_corners: 

114 # sensor_data = [output_file[f'/p/{index}'][()] for index in range(1, len(key['mask']) + 1)] 

115 # 

116 # # Combine the sensor data if using a kWaveTransducer as a sensor 

117 # if isinstance(sensor, kWaveTransducer): 

118 # sensor_data['p'] = sensor.combine_sensor_data(sensor_data['p']) 

119 

120 # # Compute the intensity outputs 

121 # if any(key.startswith(('I_avg', 'I')) for key in sensor.get('record', [])): 

122 # flags = { 

123 # 'record_I_avg': 'I_avg' in sensor['record'], 

124 # 'record_I': 'I' in sensor['record'], 

125 # 'record_p': 'p' in sensor['record'], 

126 # 'record_u_non_staggered': 'u_non_staggered' in sensor['record'] 

127 # } 

128 # kspaceFirstOrder_saveIntensity() 

129 # 

130 # # Filter the recorded time domain pressure signals using a Gaussian filter if defined 

131 # if not time_rev and 'frequency_response' in sensor: 

132 # frequency_response = sensor['frequency_response'] 

133 # sensor_data['p'] = gaussianFilter(sensor_data['p'], 1 / kgrid.dt, frequency_response[0], frequency_response[1]) 

134 # 

135 # # Assign sensor_data.p to sensor_data if sensor.record is not given 

136 # if 'record' not in sensor and not cuboid_corners: 

137 # sensor_data = sensor_data['p'] 

138 # 

139 # # Delete the input and output files 

140 # if delete_data: 

141 # os.remove(input_filename) 

142 # os.remove(output_filename) 

143 return sensor_data